fix(fused_linear_jsd): project logits in FP32 before the matmul, not after

[JeanMaximilienCadic]

fused_linear_jsd_forward casts the student/teacher logits to FP32 after the projection matmul — but the matmul already ran in bf16/fp16, so the cast just relabels a value that's already been rounded to 8/11 mantissa bits. The "compute in FP32" comment above it is effectively a no-op.

JSD's gradient is a difference of nearly-equal distributions, so that rounding cancels catastrophically. At a realistic logit spread (std≈30) I measured bf16 gradients off by up to ~20% versus a true FP32 reference — well past the tolerance CE gets away with for the same pattern.

cuBLAS already accumulates the GEMM in FP32 internally, so the fix just keeps that accumulator instead of discarding it: torch.mm(..., out_dtype=torch.float32) on torch >= 2.8 / sm_80+, with an explicit upcast fallback otherwise. Same compute and memory, correct logits.

I also fixed the transformers-level test oracle, which shared the exact same cast-after-matmul defect and was matching the kernel bug-for-bug (so CI never caught it), and added a regression test at the issue's repro shape/scale.

Out of scope (separate follow-ups, as the issue notes): the same pattern in fused_linear_distillation.py / fused_linear_ppo.py, and the grad_logits downcast before the backward matmuls.

Verification (RTX 4090, torch 2.11+cu130, triton 3.6): grad error shrinks in every case, e.g. bf16 β=0 std=30 → grad_input 13.6%→7.5%, grad_weight 15.2%→5.1%.

Fixes #1432.

[kashif]

@JeanMaximilienCadic there is already a PR that fixes this issue: #1433