Implemented vulkan backend ops for cross entropy loss and cross_entropy_loss_back - #27216
Merged
Merged
Conversation
4 tasks
Contributor
|
I'm not familiar with the cross-entropy ops, but the changes here generally look fine to me. Generally I'd prefer to handle noncontiguous and unaligned tensors when we add new ops, but I'm not sure whether it's all that important for these ops. |
0cc4m
approved these changes
Aug 26, 2026
Contributor
|
I think it's okay like this for now. |
jeffbolznv
approved these changes
Aug 26, 2026
1 task
thecodacus
pushed a commit
to thecodacus/llama.cpp
that referenced
this pull request
Sep 7, 2026
Githab-capibara
added a commit
to Githab-capibara/llama.cpp
that referenced
this pull request
Sep 10, 2026
zbrad
pushed a commit
to zbrad/llama.cpp
that referenced
this pull request
Sep 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Added vulkan support for GGML_OP_CROSS_ENTROPY_LOSS and `GGML_OP_CROSS_ENTROPY_LOSS_BACK (both are in one PR because they share the same setup except for the last pass and how the result is written where CEL_back has the full tensor with the same shape as logits, no extra kernel and scalar grad.
These used to be unsupported on vulkan and fell back to cpu. The forward path follows the cuda/SYCL one workgroup per row, numerically stable log-softmax over classes, then a sum_rows reduce into the scalar loss pattern. The backward path is a single kernel that reuses the same max + softmax setup and writes (softmax(logits) - labels) * grad / nrows
supports_op right now is limited to contiguous F32 tensors with matching logits/labels shapes (forward output must be a scalar) A
wg512pipeline variant is selected whennclasses > 1024, same threshold assoft_max.Part of #14909
Additional information
Modeled on the CUDA and SYCL cross-entropy paths, with some reuse from vulkan soft_max/sum_rows (row dispatch via ggml_vk_nrows_elements, prealloc_x temp buffer, existing pipeline_sum_rows_f32). Checked against the CPU reference in ggml_compute_forward_cross_entropy_loss_f32 and ggml_compute_forward_cross_entropy_loss_back_f32.
Also the shaders use a fixed BLOCK_SIZE workgroup reduce in shared memory, not nclasses-sized shared memory, so large-vocab cases like {30000,1,1,1} work without a separate large-class path.
Test results (AMD Radeon 780M, Windows):
Docs:
docs/ops/Vulkan.csvanddocs/ops.mdupdated for both ops on Vulkan.Requirements
YES: AI was used in the beginning to understand for research and navigation to find similar implementation. Code was all handwritten and then AI was used to review it a few times and look for edge cases and helped find some oversights/gaps which I then adressed.