Repository navigation
Add fused TP1 cross entropy with deferred backward - #550
Merged
Merged
Conversation
suachong
requested review from
RuibinCheung,
Xiaoming-AMD and
wenxie-amd
as code owners
October 10, 2026 07:30
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Destructive overwrite mode needs a safe torch.compile boundary to prevent invalid alias tracing.
1 open finding
What changed in this PR
Adds a fused TP=1 Triton cross-entropy operator with deferred gradient reconstruction.
Changes:
- Adds BF16/FP32 kernels and PyTorch API.
- Adds correctness tests, documentation, licensing, and benchmarks.
- Exports the operator publicly.
| File | Description |
|---|---|
primus_turbo/triton/cross_entropy.py |
Implements forward/backward kernels. |
primus_turbo/pytorch/ops/cross_entropy.py |
Adds validation and autograd integration. |
primus_turbo/pytorch/ops/__init__.py |
Exports the operator. |
tests/pytorch/ops/test_cross_entropy.py |
Covers correctness, layouts, and ownership. |
benchmark/ops/training/bench_cross_entropy.py |
Benchmarks Turbo and TE variants. |
docs/cross_entropy.md |
Documents usage and constraints. |
LICENSE |
Records NVIDIA-derived source licensing. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| raise ValueError("label_smoothing must be between zero and one") | ||
| if overwrite_input and not logits.is_contiguous(): | ||
| raise ValueError("overwrite_input=True requires contiguous logits") | ||
| return _CrossEntropy.apply(logits, target, ignore_index, label_smoothing, overwrite_input) |
wenxie-amd
approved these changes
Oct 11, 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.
Add a TransformerEngine-derived Triton TP=1 vocabulary cross-entropy operator with one loss/statistics forward kernel and one gradient-reconstruction backward kernel. The BF16/FP32 API supports per-token loss, ignored labels, label smoothing and supported strided inputs.
Logits are preserved by default. Optional
overwrite_input=Truereuses exclusively owned contiguous logits storage during backward. Both modes support one backward per forward; higher-order gradients are unsupported. NVIDIA Apache attribution is retained, and the operator has no runtime TransformerEngine dependency or native build changes.Validation:
[8192, 4, 128256]: installed TE 9.702 ms, Turbo preserve 6.655 ms, Turbo overwrite 5.295 ms (median, 30 repetitions after five warmups). These are operator-only measurements. The included benchmark supports reproduction against the installed TE version.ci:gpurequests GPU CI on this branch.Upstream source: TransformerEngine cross entropy. Companion integration: Primus #1245.