Skip to content

Implemented vulkan backend ops for cross entropy loss and cross_entropy_loss_back - #27216

Merged
0cc4m merged 1 commit into
ggml-org:masterfrom
PranavUttarkar:vulkan-cross-entropy-loss
Aug 26, 2026
Merged

Implemented vulkan backend ops for cross entropy loss and cross_entropy_loss_back#27216
0cc4m merged 1 commit into
ggml-org:masterfrom
PranavUttarkar:vulkan-cross-entropy-loss

Conversation

@PranavUttarkar

Copy link
Copy Markdown
Contributor

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 wg512 pipeline variant is selected when nclasses > 1024, same threshold as soft_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):

test-backend-ops.exe test    -b Vulkan0 -o CROSS_ENTROPY_LOSS       2/2 passed
test-backend-ops.exe test    -b Vulkan0 -o CROSS_ENTROPY_LOSS_BACK  2/2 passed
test-backend-ops.exe support -b Vulkan0 -o CROSS_ENTROPY_LOSS       SUPPORTED
test-backend-ops.exe support -b Vulkan0 -o CROSS_ENTROPY_LOSS_BACK  SUPPORTED
test-backend-ops.exe grad    -b Vulkan0 -o CROSS_ENTROPY_LOSS       OK ({10,5,4,3}; {30000,...} skipped as large)

Docs: docs/ops/Vulkan.csv and docs/ops.md updated 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.

@PranavUttarkar
PranavUttarkar requested a review from a team as a code owner August 17, 2026 02:30
@github-actions github-actions Bot added documentation Improvements or additions to documentation Vulkan Issues specific to the Vulkan backend ggml changes relating to the ggml tensor library for machine learning labels Aug 17, 2026
@jeffbolznv

Copy link
Copy Markdown
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

0cc4m commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

I think it's okay like this for now.

@0cc4m
0cc4m merged commit bf94216 into ggml-org:master Aug 26, 2026
27 of 30 checks passed
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation ggml changes relating to the ggml tensor library for machine learning Vulkan Issues specific to the Vulkan backend

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants