fix(ce): restore true-class gradient precision in low-precision buffers

[vaibhavjindal]

#1268 replaced the per-element tl.where(X_offsets != y, ...) selection with a post-loop read-modify-write of the true-class gradient:

tl.store(X_ptr + y, tl.load(X_ptr + y) + dxy)

X_ptr is the logits buffer, which is bf16/fp16 in mixed-precision training (FLCE computes the gradient in place). The round-trip therefore rounds softmax(x_y) / N to the buffer dtype before the -(1 - label_smoothing) / N term is subtracted.

Since dx_y = (softmax(x_y) - 1) / N cancels catastrophically as softmax(x_y) -> 1, the pre-rounding error is amplified by roughly 1 / (1 - softmax(x_y)). Measured max relative error of dx_y against an fp64 reference (bf16, V=32000, H100):

softmax(x_y)   before #1268   at #1268   torch
0.90                  0.24%      1.95%   0.43%
0.99                  0.30%     19.99%   0.46%
0.999                 0.38%    100.00%   0.42%

Confident predictions are the common case once a model has trained for a while, so this silently degrades the dominant gradient signal. It surfaced as a convergence failure in
test_mini_model[mini_gemma3_text-32-1e-05-dtype25-0.01-0.01-0.1-0.01-0.01-0.01], where 32 steps of accumulated drift push a top-k logprob past tolerance.

Recompute dx_y once in fp32 after the loop, from the already-live ori_X_y, m and d, and write it with a single store. This keeps #1268's O(1) correction (no O(V) compare/select is reintroduced) while folding in the -(1 - label_smoothing) term before the result is rounded to the buffer dtype.

Accuracy returns to the pre-#1268 values exactly (0.02% / 0.36% / 0.24% / 0.30% / 0.38%), and the speedup is retained: 3.585 ms vs 3.571 ms (this PR vs #1268) at BT=8192, V=128256 bf16 forward+backward on H100, against 3.619 ms before #1268.

The existing CE suite passed with the regression because its random logits leave softmax(x_y) near zero, so add a regression test that drives softmax(x_y) -> 1 and requires the true-class gradient to be no less accurate than torch's own low-precision backward. It fails on the #1268 kernel and passes here with ~6x headroom.

Summary

Testing Done

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence