fused_linear_jsd rounds logits to the input dtype before the documented FP32 cast, costing up to 23% gradient error in bf16

## 🐛 Describe the bug

`fused_linear_jsd_forward` intends to compute the logits in FP32, and says so:

```python
# shape: chunk_size x V
# 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)
teacher_logits_chunk = (teacher_input_chunk @ teacher_weight.t()).to(torch.float32)
```

But the matmul runs in the input dtype, so its output is rounded to 8 (bf16) or 11 (fp16) mantissa
bits *before* `.to(torch.float32)` executes. The cast is precision-inert — it changes the container,
not the contents. cuBLAS already accumulates a bf16 GEMM in FP32 internally; the kernel just throws
those bits away on the store and then casts the rounded result back up.

**Why this matters much more for JSD than for cross-entropy.** The JSD scalar is dominated by large
terms and barely moves (≤0.04% error). But `dx` is a difference of nearly-equal distributions, so the
logit rounding cancels catastrophically. In bf16 the **gradients** are off by **4–23%**.

`LigerFusedLinearJSD` is the fused path for bf16 distillation, so this lands directly on
distillation training runs.

## Reproduce

`BT=256, H=1024, V=32000`, `randn` inputs, weight scaled to a realistic logit spread. Relative error
of `max|Δ|` for the gradients, against a reference that differs from the kernel *only* in projecting
in FP32 (`x.float() @ w.float().t()`, which is exact for bf16/fp16 operands):

| dtype | beta | logit std | current `grad_input` | current `grad_weight` | with FP32 projection |
|---|---|---|---|---|---|
| bfloat16 | 0.0 (FKL) | 30 | **13.93%** | **12.20%** | 0.50% / 0.61% |
| bfloat16 | 0.5 | 10 | **22.90%** | **23.60%** | 0.38% / 0.62% |
| bfloat16 | 1.0 (RKL) | 30 | **12.35%** | **10.28%** | 0.60% / 0.71% |
| bfloat16 | 0.0 | 10 | 4.13% | 3.89% | 0.39% / 0.34% |
| float16 | 0.0 | 30 | 1.25% | 1.09% | 0.06% / 0.08% |
| float16 | 1.0 | 30 | 2.18% | 1.53% | 0.08% / 0.04% |

The error grows with logit magnitude, which is why it is worse at realistic scales than in the
existing tests.

## Why CI does not catch it

Two reasons, both in `test/transformers/test_fused_linear_jsd.py`:

1. **The oracle shares the defect.** `TorchLMHeadJSD.forward` does
`self.student_lin(student_input).to(torch.float32)` — a bf16 `nn.Linear` and then the same
inert cast. The kernel matches it *exactly*, bug for bug, so no tolerance can expose the problem.
2. **The inputs are degenerate.** Both operands come from `torch.rand` (uniform `[0, 1)`), giving
all-positive correlated vectors whose logits cluster at mean ≈ 127 with std ≈ 5. Swapping in
`randn` with a realistic logit spread is what surfaces it.

Note that `test/ops/test_fused_linear_jsd.py` already uses a genuine
`student_input.float() @ student_weight.float().T` reference, so the repo currently holds two oracles
that disagree about the intended projection precision — the `test/ops` one matches the documented
intent.

## Expected behavior

The logits should carry the FP32 accumulator the GEMM already computes, as the comment states.

`torch.mm(x, w.t(), out_dtype=torch.float32)` (torch ≥ 2.8, CUDA sm_80+) does exactly this at no
cost, and is *numerically equivalent* to casting both operands up 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.

Measured on an isolated projection GEMM at llama shapes (`BT=2048, H=4096, V=128256`):

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

End-to-end through `LigerFusedLinearJSD` (fwd+bwd, bf16, `accum_dtype` both `None` and `float32`) it
is **1.03–1.07x faster and exactly memory-neutral**. Verified that `out_dtype` is honored under
`torch.autocast(bfloat16)` and compiles under `torch.compile(fullgraph=True)`.

I have a fix ready and will open a PR against this issue.

## Out of scope here (follow-ups)

Recording these separately rather than folding them into one numerics PR:

1. **`grad_logits` is downcast to the input dtype** before the `grad_input` / `grad_weight` matmuls
(`student_logits_chunk = student_logits_chunk.to(dtype)`). After the projection fix this is the
remaining error source: `grad_input` sits at 0.38% instead of the 0.095% a fully-FP32 backward
reaches. Closing it needs either an FP32 copy of the head (+2.1 GB at llama shapes) or a
non-tensor-core FP32 GEMM, so the trade-off deserves its own discussion.

2. **The same rounding under AMP with a float32 head, which the PR's fix does not reach.**
`torch.amp.custom_fwd` is applied without `cast_inputs`, so autocast stays enabled inside the
forward and harmonizes a `bfloat16` hidden state with a `float32` `lm_head` itself. Measured at
the same shapes (`beta=0.0`, logit std 30, `autocast(bfloat16)`, float32 weights): `grad_input`
**13.18%**, `grad_weight` **11.40%** — the same magnitude as the bf16-native case.

The PR deliberately does not cover it, because unlike the same-dtype case there is no free fix:

| approach | `grad_input` | cost |
|---|---|---|
| today | 13.18% | — |
| harmonize to the autocast dtype ourselves, then `out_dtype=fp32` | 3.23% | rounds the float32 head per chunk, duplicating autocast's weight cache; still loses the head's mantissa |
| refuse the downcast, run a float32 GEMM | 0.50% | overrides the low-precision matmul the user asked autocast for, and gives up tensor cores |

Also note `aten::mm.dtype` is not on autocast's promotion list, so `torch.mm(bf16, fp32,
out_dtype=fp32)` inside an autocast region raises `RuntimeError: input dtypes must be the same`
rather than promoting — worth knowing for anyone attempting this.

3. **The same cast-after-matmul pattern in the chunked losses**, which the PR does not touch:
- `src/liger_kernel/chunked_loss/fused_linear_distillation.py:43,50` — projects with *no* FP32
cast at all.
- `src/liger_kernel/chunked_loss/fused_linear_ppo.py:47,99` — `(hidden_chunk @ weight_chunk.to(hidden.dtype).t()).float()`.

4. **Generalized JSD (`0 < beta < 1`) goes NaN for peaked distributions.** Independent of the above.
At `logit std ≈ 30`, `V=32000` the mixture `lerp(exp(log_q), exp(log_p), beta)` underflows to
exactly `0` for ~4.6M entries, so `log(m) = -inf` and `0 * -inf = nan`. The loss returns `nan`.
The mixture is formed in probability space rather than via a log-space `logsumexp`, and the
`TorchJSD` reference in `test/transformers/test_jsd.py` reproduces it identically, so this is not
specific to the Triton kernel. Repro:

```python
BT, H, V = 256, 1024, 32000
s = torch.randn(BT, H, device="cuda", dtype=torch.bfloat16)
t = torch.randn(BT, H, device="cuda", dtype=torch.bfloat16)
w = (torch.randn(V, H, device="cuda") * (30.0 / H**0.5)).to(torch.bfloat16)
wt = (torch.randn(V, H, device="cuda") * (30.0 / H**0.5)).to(torch.bfloat16)
LigerFusedLinearJSD(jsd_beta=0.5)(s.requires_grad_(), w.requires_grad_(), t, wt, None) # -> nan
```

## Environment report

```
Torch: 2.13.0+cu130
GPU: NVIDIA GB10 (sm_121)
Liger-Kernel: main (0.8.2)
```

Start at the fused_linear_jsd_forward entry point and compare its projection with the documented FP32 behavior. Run test/ops/test_fused_linear_jsd.py and test/transformers/test_fused_linear_jsd.py, using the genuine FP32 projection reference already described in the issue. Done means bf16 and fp16 logits preserve the FP32 GEMM accumulator without changing the out-of-scope cases.

Written by the indexing model from the issue text.