Refactor Triton cross-entropy executor: cleanup | Reduce activation memory in cross-entropy backward

[oonyshch]

Summary

Hi!
I'd like to make my initial contribution to Thunder. I've been reading through the Triton cross-entropy executor (triton_crossentropy_impl.py, triton_crossentropy.py, test_triton_ce.py). Found dead code, unnecessary indirection, misleading naming. I'd like to clean it all up in a single PR.
Then follow up with a performance improvement PR for the same scope next, described below.

Findings

Dead code, unnecessary indirection in triton_crossentropy_impl.py:

What	Problem
TritonDtype enum + _TORCH2DTYPE + _DTYPE2TRITON	Three objects to do a torch.dtype → tl.dtype mapping: and that can be a single dict
FORWARD_NUM_STAGES = 1	Marked "Temporarily borrowed from openai/triton", equals 1, where backward already hardcodes 1 directly
buffer_dtype = BUFFER_DTYPE at the start of every kernel	Alias for a tl.constexpr that can be used directly
Commented-out assert and empty_like for indices	Leftover debug code
Extra if triton is None: return False in cross_entropy_checker	Unreachable, because module-level assert TRITON_AVAILABLE guarantees Triton is present
buffer_dtype = None + if buffer_dtype is None: in CrossEntropy.forward	Always-None init + always-true branch
import thunder.torch as ltorch at line ~620	Module-level import buried below - Mostly across repo it is on top, should be on top
# for start_n in range(0, ...): # need to change this	Stale comment

Misleading variable name

The kernel variable log_softmax computes log(Σexp) − xᵢ, which is the negative log-softmax. The Python side correctly calls it neg_logprobs, but the kernel name inverts the sign semantics. The backward confirms this by loading it with -tl.load(...).

More minor issues

• Docstring typo: "indcies" → "indices", "probabilites" → "probabilities"
• Stale comment in backward: "write result in-place in PROBS": writes to DIN, not PROBS
• min_triton_version = "2.1" is duplicated independently in triton_crossentropy_impl.py and test_triton_ce.py.

Proposed changes (one PR)

All changes are in triton_crossentropy_impl.py:

1. Replace TritonDtype enum + _TORCH2DTYPE + _DTYPE2TRITON with a single _TORCH2TRITON_DTYPE dict
2. Inline FORWARD_NUM_STAGES = 1 → use 1 directly in autotune configs
3. Remove buffer_dtype = BUFFER_DTYPE aliases and use BUFFER_DTYPE directly in kernels
4. Remove commented-out assert and empty_like
5. Remove unreachable if triton is None guard in checker
6. Remove always-true buffer_dtype = None / if buffer_dtype is None: wrapper
7. Move import thunder.torch as ltorch to top-level imports
8. Rename kernel variable log_softmax → neg_log_softmax
9. Remove stale # for start_n ... comment
10. Fix docstring typos ("indcies", "probabilites")
11. Fix stale backward comment ("PROBS" → "DIN")

Future optimization (another PR)

Currently CrossEntropy.forward saves the full neg_logprobs tensor of shape (B, N) , which is batch size × vocabulary size , for use in the backward pass:

neg_logprobs = torch.empty_like(logits, dtype=buffer_dtype, device=device)  # (B, N)
ctx.save_for_backward(neg_logprobs, indices, weights_buffer)

For a 128k-vocab LLM with B=2048 in float32, this costs 2048 × 128000 × 4 bytes ≈ 1 GB.

However, the backward only ever does probs = tl.exp(probs) on it to recover the softmax. Since:

$$\text{softmax}_i = \exp(x_i - \max - \log\sum\exp)$$

the softmax can be recomputed from just max and log_sum_exp per row (already computed in the forward's online softmax loop) plus the original logits. That reduces saved activations from O(B×N) complexity to just O(B), because two scalars per row instead of the full row.

The trade-off is an extra read of the logits in the backward, but for large-vocab models the memory savings far outweigh this. This is the approach used by Liger Kernel's cross-entropy.
So I'd like to tackle this as a follow-up PR.

[IvanYashchuk]

Hi, thanks for the thorough write-up! You clearly spent time reading through the code, and the findings are accurate. The TritonDtype enum chain, the sign mismatch in the naming, the always-None buffer_dtype init: all good catches. The memory optimization analysis for the backward pass is sound too. I want to be upfront: our team's priorities have shifted toward torch.compile and upstream PyTorch work, and Thunder is not receiving active development from my team right now. I wouldn't want you to invest time on a PR that won't get the review attention it deserves. It wouldn't be fair to your effort.
That said, the quality of this analysis speaks for itself. If you're interested in the kind of compiler and optimization work our team focuses on, take a look at what's happening in the torch.compile space. Happy to chat more about that separately.