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

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:

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.

Start by reading triton_crossentropy_impl.py, then inspect test_triton_ce.py for the existing Triton cross-entropy coverage and version requirement. Apply the listed cleanup, naming, import, typo, and comment changes in the implementation file, keeping activation-memory reduction as a separate follow-up, and run the Triton cross-entropy tests.