perf(ce): single-shot dx_y gradient correction

[justinhh4]

🔗 Stacked on #1266 (exp2) and #1267 (num_warps). Please review/merge after those. This PR's own change is the dx_y commit (perf(ce): single post-loop dx_y correction) — once the other two merge, this diff collapses to just that.

PR3: CE exp2 + num_warps + dx_y — B200 + H100 benchmarks

Optimization measured against the frozen original baseline (v0_base, == upstream HEAD). Correctness: full CE suite (177 cases, bf16+fp32, weight/softcap/smoothing/z_loss/ignore_index combos) passes.

What & why

The CE gradient subtracts 1 only at the true-class index y. Instead of a per-element tl.where(X_offsets != y, …) across the whole vocab (in both pass-2 paths), store the plain softmax gradient and apply one scalar correction at index y after the loop.

• Why: the per-element compare+select ran O(V) times per row just to touch a single element; the post-loop store makes it O(1). A tl.debug_barrier() orders the loop's in-place writes before the read-back at y.
• Real removal, not predication: no conditional remains in the hot loop, so unlike loop-splitting (which Triton predicates) this genuinely deletes instructions. The correction carries the same weight / reduction / softcap factors the per-element path applied.

Headline shape (BT=8192, V=128256, full fwd+bwd)

GPU	dtype	baseline	this PR	latency	throughput
H100	bf16	2.583 ms	2.198 ms	-14.9%	+17.5%
H100	fp32	4.489 ms	4.488 ms	-0.0%	+0.0%
B200	bf16	1.847 ms	1.193 ms	-35.4%	+54.9%
B200	fp32	2.085 ms	1.832 ms	-12.1%	+13.8%

Across context length (V=128256, BT 1024–65536)

GPU	dtype	min	avg	max
H100	bf16	-9%	-14%	-17%
H100	fp32	+0%	-0%	-0%
B200	bf16	-16%	-32%	-40%
B200	fp32	-5%	-11%	-13%

Figures

Latency vs context length and vs vocab, baseline vs this PR, on H100 and B200.

[vaibhavjindal]

@justinhh4 since #1266 and #1267 are merged, we can try merging this now. Can you update the PR description reporting just the benefits from single shot dx_y improvement? This will help us evaluate the individual impact of this change.