vulkan backend ops: implemented GATED_LINEAR_ATTN (#25601)
* vulkan : add GATED_LINEAR_ATTN op * docs : update Vulkan ops * vulkan : remove unused GLA spec constant * Updated ops.md * ops.md update
This commit is contained in:
@@ -1026,6 +1026,7 @@ struct vk_device_struct {
|
||||
vk_pipeline pipeline_pool2d_f32;
|
||||
vk_pipeline pipeline_rwkv_wkv6_f32;
|
||||
vk_pipeline pipeline_rwkv_wkv7_f32;
|
||||
vk_pipeline pipeline_gated_linear_attn_f32;
|
||||
// [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128
|
||||
vk_pipeline pipeline_gated_delta_net[4][2];
|
||||
vk_pipeline pipeline_ssm_scan_f32_d128;
|
||||
@@ -1747,6 +1748,13 @@ struct vk_op_rwkv_wkv7_push_constants {
|
||||
uint32_t C;
|
||||
uint32_t H;
|
||||
};
|
||||
struct vk_op_gated_linear_attn_push_constants {
|
||||
uint32_t B;
|
||||
uint32_t T;
|
||||
uint32_t C;
|
||||
uint32_t H;
|
||||
float scale;
|
||||
};
|
||||
struct vk_op_gated_delta_net_push_constants {
|
||||
uint32_t H;
|
||||
uint32_t n_tokens;
|
||||
@@ -5665,6 +5673,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
|
||||
|
||||
ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_f32, "rwkv_wkv7_f32", rwkv_wkv7_f32_len, rwkv_wkv7_f32_data, "main", 8, sizeof(vk_op_rwkv_wkv7_push_constants), {1, 1, 1}, {device->subgroup_size}, 1);
|
||||
|
||||
ggml_vk_create_pipeline(device, device->pipeline_gated_linear_attn_f32, "gated_linear_attn_f32", gated_linear_attn_f32_len, gated_linear_attn_f32_data, "main", 6, sizeof(vk_op_gated_linear_attn_push_constants), {1, 1, 1}, {}, 1);
|
||||
|
||||
{
|
||||
const uint32_t gdn_sizes[] = {16, 32, 64, 128};
|
||||
const char * gdn_names[][2] = {
|
||||
@@ -11392,6 +11402,11 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
|
||||
return ctx->device->pipeline_rwkv_wkv7_f32;
|
||||
}
|
||||
return nullptr;
|
||||
case GGML_OP_GATED_LINEAR_ATTN:
|
||||
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
|
||||
return ctx->device->pipeline_gated_linear_attn_f32;
|
||||
}
|
||||
return nullptr;
|
||||
case GGML_OP_GATED_DELTA_NET:
|
||||
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
|
||||
const uint32_t S_v = dst->src[2]->ne[0];
|
||||
@@ -12422,6 +12437,41 @@ static void ggml_vk_rwkv_wkv7(ggml_backend_vk_context * ctx, vk_context& subctx,
|
||||
);
|
||||
}
|
||||
|
||||
static void ggml_vk_gated_linear_attn(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
|
||||
const size_t seq_length = dst->src[0]->ne[2];
|
||||
const size_t n_embed = dst->ne[0];
|
||||
const size_t n_heads = dst->src[0]->ne[1];
|
||||
const size_t n_seqs = dst->src[4]->ne[1];
|
||||
|
||||
float scale;
|
||||
memcpy(&scale, dst->op_params, sizeof(float));
|
||||
|
||||
GGML_ASSERT(dst->buffer != nullptr);
|
||||
|
||||
vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, dst->src[0], dst->src[1], dst->src[2], dst, dst->op);
|
||||
GGML_ASSERT(pipeline != nullptr);
|
||||
|
||||
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
|
||||
|
||||
vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst);
|
||||
vk_subbuffer src_buf[5] = {};
|
||||
for (int i = 0; i < 5; i++) {
|
||||
src_buf[i] = ggml_vk_tensor_subbuffer(ctx, dst->src[i]);
|
||||
}
|
||||
|
||||
const vk_op_gated_linear_attn_push_constants pc = {
|
||||
(uint32_t)n_seqs,
|
||||
(uint32_t)seq_length,
|
||||
(uint32_t)n_embed,
|
||||
(uint32_t)n_heads,
|
||||
scale,
|
||||
};
|
||||
|
||||
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
|
||||
{src_buf[0], src_buf[1], src_buf[2], src_buf[3], src_buf[4], dst_buf},
|
||||
pc, { (uint32_t)(n_seqs * n_heads), 1, 1 });
|
||||
}
|
||||
|
||||
static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
|
||||
const ggml_tensor * src_q = dst->src[0];
|
||||
const ggml_tensor * src_v = dst->src[2];
|
||||
@@ -15421,6 +15471,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
|
||||
|
||||
break;
|
||||
|
||||
case GGML_OP_GATED_LINEAR_ATTN:
|
||||
ggml_vk_gated_linear_attn(ctx, compute_ctx, node);
|
||||
|
||||
break;
|
||||
|
||||
case GGML_OP_GATED_DELTA_NET:
|
||||
ggml_vk_gated_delta_net(ctx, compute_ctx, node);
|
||||
|
||||
@@ -18128,6 +18183,9 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
|
||||
case GGML_OP_RWKV_WKV6:
|
||||
case GGML_OP_RWKV_WKV7:
|
||||
return true; // all inputs are contiguous, see ggml.c
|
||||
case GGML_OP_GATED_LINEAR_ATTN:
|
||||
// the shader block size is hardcoded to head_size 64
|
||||
return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && op->src[0]->ne[0] == 64;
|
||||
case GGML_OP_GATED_DELTA_NET:
|
||||
{
|
||||
const uint32_t S_v = op->src[2]->ne[0];
|
||||
@@ -19117,6 +19175,10 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
|
||||
} else if (tensor->op == GGML_OP_RWKV_WKV7) {
|
||||
tensor_clone = ggml_rwkv_wkv7(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3],
|
||||
src_clone[4], src_clone[5], src_clone[6]);
|
||||
} else if (tensor->op == GGML_OP_GATED_LINEAR_ATTN) {
|
||||
const float * op_params = (const float *)tensor->op_params;
|
||||
tensor_clone = ggml_gated_linear_attn(ggml_ctx, src_clone[0], src_clone[1],
|
||||
src_clone[2], src_clone[3], src_clone[4], op_params[0]);
|
||||
} else if (tensor->op == GGML_OP_GATED_DELTA_NET) {
|
||||
tensor_clone = ggml_gated_delta_net(ggml_ctx, src_clone[0], src_clone[1],
|
||||
src_clone[2], src_clone[3], src_clone[4], src_clone[5],
|
||||
|
||||
Reference in New Issue
Block a user