Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Implemented vulkan cross_entropy_loss and cross_entropy_loss_back (#27216)

Browse Source
This commit is contained in:
Pranav Uttarkar
2026-08-26 16:49:32 +02:00
committed by GitHub
parent d0132a680a
commit bf94216469
6 changed files with 299 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+138
View File
@@ -1042,6 +1042,8 @@ struct vk_device_struct {
vk_pipeline pipeline_argsort_large_f32[num_argsort_pipelines];
vk_pipeline pipeline_topk_f32[num_topk_pipelines];
vk_pipeline pipeline_sum_rows_f32;
vk_pipeline pipeline_cross_entropy_loss_f32, pipeline_cross_entropy_loss_f32_wg512;
vk_pipeline pipeline_cross_entropy_loss_back_f32, pipeline_cross_entropy_loss_back_f32_wg512;
vk_pipeline pipeline_fwht_f32[4];
vk_pipeline pipeline_cumsum_f32;
vk_pipeline pipeline_cumsum_small_f32;
@@ -5758,6 +5760,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_f32, "cross_entropy_loss_f32", cross_entropy_loss_f32_len, cross_entropy_loss_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_f32_wg512, "cross_entropy_loss_f32_wg512", cross_entropy_loss_f32_len, cross_entropy_loss_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, { 512 }, 1);
ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_back_f32, "cross_entropy_loss_back_f32", cross_entropy_loss_back_f32_len, cross_entropy_loss_back_f32_data, "main", 4, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_back_f32_wg512, "cross_entropy_loss_back_f32_wg512", cross_entropy_loss_back_f32_len, cross_entropy_loss_back_f32_data, "main", 4, sizeof(vk_op_push_constants), {1, 1, 1}, { 512 }, 1);
// Intel Windows driver in range [32.0.101.8509, 32.0.101.8860) will crash when using fwht kernels so we gate that here
const bool can_use_fwht = device->driver_id != vk::DriverId::eIntelProprietaryWindows ||
!ggml_vk_intel_windows_driver_in_range(device->properties.driverVersion, 101, 8509, 101, 8860);
@@ -11577,6 +11583,17 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
return ctx->device->pipeline_sum_rows_f32;
}
return nullptr;
case GGML_OP_CROSS_ENTROPY_LOSS:
if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
return src0->ne[0] > 1024 ? ctx->device->pipeline_cross_entropy_loss_f32_wg512 : ctx->device->pipeline_cross_entropy_loss_f32;
}
return nullptr;
case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
// src0 is the scalar grad; src1 is logits
if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && src2 && src2->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
return src1->ne[0] > 1024 ? ctx->device->pipeline_cross_entropy_loss_back_f32_wg512 : ctx->device->pipeline_cross_entropy_loss_back_f32;
}
return nullptr;
case GGML_OP_CUMSUM:
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
if (src0->ne[0] <= 512) {
@@ -13942,6 +13959,103 @@ static void ggml_vk_cumsum(ggml_backend_vk_context * ctx, vk_context& subctx, co
ctx->prealloc_split_k_need_sync = true;
}
static std::array<uint32_t, 3> ggml_vk_nrows_elements(uint32_t nr) {
if (nr > 262144) {
return { 512, 512, CEIL_DIV(nr, 262144) };
}
if (nr > 512) {
return { 512, CEIL_DIV(nr, 512), 1 };
}
return { nr, 1, 1 };
}
static void ggml_vk_cross_entropy_loss(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
GGML_ASSERT(src0->type == GGML_TYPE_F32);
GGML_ASSERT(src1->type == GGML_TYPE_F32);
GGML_ASSERT(dst->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_is_contiguous(src0));
GGML_ASSERT(ggml_is_contiguous(src1));
GGML_ASSERT(ggml_is_contiguous(dst));
GGML_ASSERT(ggml_are_same_shape(src0, src1));
GGML_ASSERT(ggml_is_scalar(dst));
const uint32_t nclasses = (uint32_t)src0->ne[0];
const uint32_t nrows = (uint32_t)ggml_nrows(src0);
vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, nullptr, dst, GGML_OP_CROSS_ENTROPY_LOSS);
GGML_ASSERT(pipeline != nullptr);
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_sum_rows_f32, 1);
vk_subbuffer src0_buf = ggml_vk_tensor_subbuffer(ctx, src0);
vk_subbuffer src1_buf = ggml_vk_tensor_subbuffer(ctx, src1);
vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true);
const vk_op_push_constants pc = { nclasses, nrows, 0.0f, 0.0f, 0.0f, 0.0f };
const size_t tmp_size = (size_t)nrows * sizeof(float);
if (ctx->prealloc_size_x < tmp_size) {
ctx->prealloc_size_x = tmp_size;
ggml_vk_preallocate_buffers(ctx, subctx);
}
if (ctx->prealloc_x_need_sync) {
ggml_vk_sync_buffers(ctx, subctx);
}
vk_subbuffer tmp_buf = { ctx->prealloc_x, 0, tmp_size };
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src0_buf, src1_buf, tmp_buf }, pc, ggml_vk_nrows_elements(nrows));
ggml_vk_sync_buffers(ctx, subctx);
vk_op_sum_rows_push_constants sp = {};
sp.n_cols = nrows;
sp.ne01 = 1;
sp.ne02 = 1;
sp.weight = 1.0f;
init_pushconst_fastdiv(sp);
sp.misalign_offsets = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_sum_rows_f32, { tmp_buf, dst_buf }, sp, { 1, 1, 1 });
ctx->prealloc_x_need_sync = true;
}
static void ggml_vk_cross_entropy_loss_back(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
const ggml_tensor * grad = dst->src[0];
const ggml_tensor * logits = dst->src[1];
const ggml_tensor * labels = dst->src[2];
GGML_ASSERT(grad->type == GGML_TYPE_F32);
GGML_ASSERT(logits->type == GGML_TYPE_F32);
GGML_ASSERT(labels->type == GGML_TYPE_F32);
GGML_ASSERT(dst->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_is_scalar(grad));
GGML_ASSERT(ggml_is_contiguous(grad));
GGML_ASSERT(ggml_is_contiguous(logits));
GGML_ASSERT(ggml_is_contiguous(labels));
GGML_ASSERT(ggml_is_contiguous(dst));
GGML_ASSERT(ggml_are_same_shape(logits, labels));
GGML_ASSERT(ggml_are_same_shape(logits, dst));
const uint32_t nclasses = (uint32_t)logits->ne[0];
const uint32_t nrows = (uint32_t)ggml_nrows(logits);
vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, grad, logits, labels, dst, GGML_OP_CROSS_ENTROPY_LOSS_BACK);
GGML_ASSERT(pipeline != nullptr);
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
vk_subbuffer grad_buf = ggml_vk_tensor_subbuffer(ctx, grad);
vk_subbuffer logits_buf = ggml_vk_tensor_subbuffer(ctx, logits);
vk_subbuffer labels_buf = ggml_vk_tensor_subbuffer(ctx, labels);
vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst);
const vk_op_push_constants pc = { nclasses, nrows, 0.0f, 0.0f, 0.0f, 0.0f };
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { grad_buf, logits_buf, labels_buf, dst_buf }, pc, ggml_vk_nrows_elements(nrows));
}
static void ggml_vk_argmax(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_ARGMAX, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], 0.0f, 0.0f, 0.0f, 0.0f });
}
@@ -15687,6 +15801,14 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
case GGML_OP_ARGMAX:
ggml_vk_argmax(ctx, compute_ctx, src0, node);
break;
case GGML_OP_CROSS_ENTROPY_LOSS:
ggml_vk_cross_entropy_loss(ctx, compute_ctx, node);
break;
case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
ggml_vk_cross_entropy_loss_back(ctx, compute_ctx, node);
break;
case GGML_OP_COUNT_EQUAL:
ggml_vk_count_equal(ctx, compute_ctx, src0, src1, node);
@@ -18511,6 +18633,18 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
}
case GGML_OP_ARGMAX:
return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
case GGML_OP_CROSS_ENTROPY_LOSS:
return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32
&& ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32
&& ggml_are_same_shape(op->src[0], op->src[1])
&& ggml_is_contiguous(op) && ggml_is_scalar(op) && op->type == GGML_TYPE_F32;
case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32 && ggml_is_scalar(op->src[0])
&& ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32
&& ggml_is_contiguous(op->src[2]) && op->src[2]->type == GGML_TYPE_F32
&& ggml_are_same_shape(op->src[1], op->src[2])
&& ggml_are_same_shape(op->src[1], op)
&& ggml_is_contiguous(op) && op->type == GGML_TYPE_F32;
case GGML_OP_COUNT_EQUAL:
return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_I32
&& ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_I32;
@@ -19437,6 +19571,10 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
tensor_clone = ggml_mean(ggml_ctx, src_clone[0]);
} else if (tensor->op == GGML_OP_ARGMAX) {
tensor_clone = ggml_argmax(ggml_ctx, src_clone[0]);
} else if (tensor->op == GGML_OP_CROSS_ENTROPY_LOSS) {
tensor_clone = ggml_cross_entropy_loss(ggml_ctx, src_clone[0], src_clone[1]);
} else if (tensor->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK) {
tensor_clone = ggml_cross_entropy_loss_back(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]);
} else if (tensor->op == GGML_OP_COUNT_EQUAL) {
tensor_clone = ggml_count_equal(ggml_ctx, src_clone[0], src_clone[1]);
} else if (tensor->op == GGML_OP_SOLVE_TRI) {