Skip to content

[CUDA] Ce fused kernel - #3947

Draft
nastya236 wants to merge 5 commits into
mainfrom
ce-fused-kernel
Draft

[CUDA] Ce fused kernel#3947
nastya236 wants to merge 5 commits into
mainfrom
ce-fused-kernel

Conversation

@nastya236

@nastya236 nastya236 commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator

Fused backward + forward kernels for cross entropy.

There is an alternative to make this a custom kernel and put it into pretraining pipeline (as in mlx-lm for kl and js losses). But after given it a thought, I am leaning towards having a fast primitive for it. I feel that ce loss is an essential part of next token prediction and it deserves to be a fast primitive:)

forward: compute loss = logsumexp(x) - x_n -> logits and loss are saved for backward
backward: compute exp(loss + x_n) - onehot, donate logits buffer for dx (we can do it in most cases)

logits = model(x).astype(mx.float32) <- now after forward we don't need to cast to float32 and it saves crazy amount of memory: for 4B pretraining with input.shape = (2, 8192) and vocab = 128256 it results in more than 17GB decrease of peak memory!

I am happy to close this PR and make it a custom kernel if you think it is a preferable option.

@nastya236 nastya236 changed the title Ce fused kernel [CUDA] Ce fused kernel Jul 30, 2026
@nastya236
nastya236 marked this pull request as draft July 30, 2026 13:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant