Conversation
CrossEntropy and CrossEntropyVJP had a fused CUDA kernel but threw "Metal cross entropy NYI" on Metal, where use_fallback always returned true and the op decomposed into logsumexp + take_along_axis. Port the CUDA kernel to Metal, following the logsumexp.h idiom: a block variant for rows up to 4096 and a looped variant above that, both reducing with simd_max/simd_sum and a cross-simdgroup pass through threadgroup memory. The accumulation is in float32 regardless of input type, so half precision inputs no longer materialize a float32 copy of the logits. The vjp writes the gradient into the logits buffer when that buffer is donatable. Each thread reads the target score into a register and the threadgroup barriers before any of them overwrite the row. Fusing also avoids a catastrophic cancellation the decomposed path suffers: forming (max - x_target) before adding log(normalizer) keeps the loss when the logits are large, where computing the full logsumexp first rounds it away. On an M2 Max, forward + backward is 2.4-6.1x faster and peak memory for 4096 x 128256 bfloat16 drops from 4.89 GiB to 1.96 GiB. Forward alone is 2.2-4.1x faster in half precision and roughly at parity in float32, where the fallback needs no cast. Validated against the decomposed fallback for float32/float16/bfloat16 across vocabularies spanning the block/looped boundary, degenerate rows (+inf, fully masked), non-contiguous logits and targets, donated buffers, and with both MLX_METAL_JIT on and off.
The use_fast gate checked mx.cuda.is_available(), so the fused primitive was unreachable from nn.losses.cross_entropy on Metal even once the Metal kernel existed. Both GPU backends implement it now, so gate on the device rather than on the backend. The fast branch does not apply the max shift the fallback branch does, because the fused kernel shifts by the row max internally. The Metal kernel has that property, so the existing shift invariance tests hold on Metal: the loss stays log(n) for equal logits from 1e0 to 1e20, and matches the cpu at every offset. Add a test that the two devices agree. In float32 they agree exactly. In half precision they are allowed to differ, because the fused kernel accumulates in float32 while the fallback reduces in the dtype of the logits, so the test asserts the stronger property instead: the fused result is never further from the float32 answer. It is a regression guard rather than a demonstration of this change, since routing to a more accurate kernel cannot make an equivalence test fail. Also correct the docstring, which said Metal reduces in the dtype of the logits, and the dtype in the half precision example, which is the dtype of the logits rather than float32.
Collaborator
|
I believe this pull request is a duplication of #4520, is not it? |
nastya236
added
the
await response
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
2 participants
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.
Summary
Validation