fix: request float32 logits from the fused JSD projection GEMM

[KyleMylonakisProtopia]

Summary

fused_linear_jsd_forward documents that it computes the logits in FP32:

# For anything starting from logits to the final JSD loss, we do computation
# in FP32 to avoid losing numerical stability.
student_logits_chunk = (student_input_chunk @ student_weight.t()).to(torch.float32)

It does not. The matmul runs in the input dtype, so the logits are rounded to 8 (bf16) or 11 (fp16)
mantissa bits before the cast executes — the cast changes the container, not the contents. cuBLAS
already accumulates a bf16 GEMM in FP32; the kernel discards those bits on the store, then casts the
rounded result back up.

The JSD scalar hides this (≤0.04% error) because it is dominated by large terms. dx does not: it is
a difference of nearly-equal distributions, so the rounding cancels catastrophically and the
gradients are off by 4–23% in bf16. LigerFusedLinearJSD is the fused path for bf16 distillation,
so this lands on distillation training runs.

This asks the projection GEMM for an FP32 output instead of casting after the fact.

Closes #1432

Why out_dtype rather than upcasting the operands

torch.mm(x, w.t(), out_dtype=torch.float32) is numerically equivalent to casting both operands to
FP32 and running an FP32 GEMM: a bf16×bf16 or fp16×fp16 product is exactly representable in FP32
(8+8 and 11+11 mantissa bits both fit in 24) and the accumulation is FP32 either way. Unlike an
explicit upcast it keeps the low-precision tensor cores and needs no FP32 copy of the head.

Isolated projection GEMM, BT=2048, H=4096, V=128256:

projection	time	peak transient	max abs logit error vs FP64
(x @ w.t()).to(fp32) (before)	29.4 ms	1504 MiB	1.0e+00
torch.mm(..., out_dtype=fp32) (this PR)	25.2 ms	1002 MiB	2.1e-03
upcast both operands to fp32	196.2 ms	1802 MiB	8.7e-04

Accuracy

Relative error of max|Δ| for the gradients against a reference differing from the kernel only in
projecting in FP32. BT=256, H=1024, V=32000, randn inputs, weight scaled to the stated logit std:

dtype	beta	logit std	grad_input before → after	grad_weight before → after
bfloat16	0.0 (FKL)	30	13.93% → 0.50%	12.20% → 0.61%
bfloat16	0.5	10	22.90% → 0.38%	23.60% → 0.62%
bfloat16	1.0 (RKL)	30	12.35% → 0.60%	10.28% → 0.71%
float16	0.0	30	1.25% → 0.06%	1.09% → 0.08%
float16	1.0	30	2.18% → 0.08%	1.53% → 0.04%

Performance

End-to-end fwd+bwd through LigerFusedLinearJSD, bf16, mean of 10 iterations after 3 warmups:

BT	H	V	accum_dtype	before	after	speedup	peak memory
4096	4096	128256	None	662.21 ms	640.30 ms	1.034x	13186 MiB → 13186 MiB
2048	4096	32000	None	82.05 ms	79.36 ms	1.034x	3096 MiB → 3096 MiB
1024	2048	32000	None	31.61 ms	30.47 ms	1.038x	1548 MiB → 1548 MiB
4096	4096	128256	float32	705.77 ms	681.43 ms	1.036x	14689 MiB → 14689 MiB
2048	4096	32000	float32	89.88 ms	86.47 ms	1.039x	3346 MiB → 3346 MiB
1024	2048	32000	float32	35.96 ms	33.60 ms	1.070x	1673 MiB → 1673 MiB

Faster, and memory-neutral end-to-end. (The isolated GEMM saves the bf16 logits buffer, but the
caching allocator absorbs it at the kernel level — so no memory win is claimed here.)

This changes bf16/fp16 numerics by default

Deliberately, and with no new flag: there is no trade-off to expose, since the FP32 projection is
more accurate, faster, and memory-neutral. Unlike accum_dtype — which is opt-in because FP32
accumulation costs a real V x H x 4 buffer — this costs nothing, so gating it would leave every
user on the inaccurate path. The kernel's own comment already promises this behavior.

Compatibility:

• All existing tolerances pass unchanged, including the bf16 test_correctness,
  test_correctness_with_ignore_index, test_correctness_functional, test_correctness_all_ignored
  and test_amp cases, which compare against the bf16-projection TorchLMHeadJSD oracle.
• out_dtype is honored under torch.autocast(bfloat16) (verified) and compiles under
  torch.compile(fullgraph=True) (verified).
• Platforms without out_dtype support (torch < 2.8, non-CUDA, sm < 80) keep today's expression, so
  there is no regression and no new OOM risk there. This mirrors how the accum_dtype grad-weight
  path in this same function already guards its torch.addmm(..., out_dtype=...) fast path.
• grad_input, the accum_dtype grad-weight path, fused_linear_jsd_backward, the
  dispatch("jsd_loss_and_grad", ...) inner kernel and the autograd signature are all untouched, so
  the cuTile / CuTeDSL JSD backends are unaffected.

The existing _ADDMM_SUPPORTS_OUT_DTYPE gate is renamed _MM_SUPPORTS_OUT_DTYPE since it now guards
both torch.mm and torch.addmm — same torch 2.8 feature, same value. Happy to drop the rename if
you would rather keep the diff to the projection alone.

What this does not fix: AMP with a float32 head

The guard requires the operands to share a dtype, so it deliberately skips the autocast case.
torch.amp.custom_fwd is applied without cast_inputs, so autocast stays enabled inside the forward
and harmonizes a bf16 hidden state with an fp32 lm_head itself. That path has the same magnitude of
error (grad_input 13.18%, grad_weight 11.40% at the shapes above) and this PR does not improve it.

That is a scope decision, not an oversight: unlike the same-dtype case there is no free fix. Rounding
the fp32 head ourselves once per chunk duplicates autocast's weight cache and still loses the head's
mantissa (13.18% → 3.23%); refusing the downcast for an fp32 GEMM gets 13.18% → 0.50% but overrides
the low-precision matmul the user explicitly asked autocast for and gives up tensor cores. Both are
trade-offs worth their own discussion, so they are recorded in #1432 rather than folded in here.
(Also worth knowing: aten::mm.dtype is not on autocast's promotion list, so passing out_dtype
inside an autocast region raises rather than promoting.) test_amp is unaffected and still passes.

Testing

New regression test test_logits_projection_is_not_rounded_to_input_dtype, parametrized over
dtype x beta. It compares against an FP32-projection oracle rather than the existing
TorchLMHeadJSD, which cannot detect this defect because it performs the same inert cast (a bf16
nn.Linear followed by .to(torch.float32)) and so matches the buggy kernel exactly. The existing
cases also draw both operands from torch.rand, clustering the logits at mean ≈ 127 / std ≈ 5;
randn at a realistic spread is what exposes the rounding.

Verified as a real regression test: with the kernel change reverted, all 5 active parametrizations
fail at 1.25%–22.90%; with it, all pass under 1%. TorchLMHeadJSD is intentionally left alone — it
still passes and legitimately documents parity with an eager bf16 lm_head.

The float16, beta=0.5 combination is explicitly skipped with its reason: there the reference
gradient peaks at ~4e-6, below float16's smallest normal (6.1e-5), so quantization of the output
buffer dominates and the projection cannot be isolated. The FP32 oracle itself misses the 1% bar
there, so asserting it would be testing float16's exponent range, not this change.

make checkstyle is clean. make test:

	passed	failed	skipped	xfailed
this branch	4457	2	980	15
main (control, same machine)	4452	2	979	15

The +5 passed is exactly this PR's new parametrizations. The 2 failures are pre-existing on
main: I ran the same suite on an unmodified checkout and it fails the identical two tests,
test/transformers/test_grpo_loss.py::test_grpo_loss_vs_trl[2-128-1000-0.04-{False,True}-luspo-sequence-1.5],
with byte-identical NaN gradient tensors. They also pass in isolation and at file scope (104 passed,
17 skipped), so they are order-dependent within the full run. test_grpo_loss.py and
chunked_loss/grpo_loss.py contain no reference to JSD or fused_linear_jsd. Happy to file that
separately if it is not already tracked.

Every JSD-adjacent suite is green on this branch:

pytest test/transformers/test_fused_linear_jsd.py test/ops/test_fused_linear_jsd.py \
       test/ops/test_jsd.py test/transformers/test_jsd.py \
       test/chunked_loss/test_jsd_loss.py test/ops/test_jsd_chunk_ignore_regression.py
# 245 passed, 34 skipped

Tolerances in test/ops/test_fused_linear_jsd.py are deliberately left alone: that file
parametrizes over every registered backend, whose inner JSD kernels differ.

[KyleMylonakisProtopia]

Any chance of getting this in? It would fix a significant loss of precision and be faster than the current implementation for virtually no change in memory overhead? Seems like a no-brainer.