fix: request float32 logits from the fused JSD projection GEMM (#1433)

## Summary

`fused_linear_jsd_forward` documents that it computes the logits in
FP32:

```python
# 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.

---------

Co-authored-by: Claude Opus 5 (1M context) <[email protected]>
Co-authored-by: Vaibhav Jindal <[email protected]>

src/liger_kernel/ops/fused_linear_jsd.py

@@ -22,7 +22,8 @@ from liger_kernel.utils import infer_device
 # The optimal maximum block size depends on your hardware, your kernel, and your dtype
 MAX_FUSED_SIZE = 4096 if infer_device() == "xpu" else 65536 // 2
 _TORCH_VERSION = Version(torch.__version__.split("+")[0])
-_ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0")
+# torch 2.8 added ``out_dtype=`` to both ``torch.mm`` and ``torch.addmm``.
+_MM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0")
 
 
 def _max_capability(device=None) -> int:
@@ -45,6 +46,40 @@ def _max_capability(device=None) -> int:
     return 0
 
 
+def _project_to_float32(inputs, weight):
+    """Project ``inputs @ weight.t()`` into float32 logits without rounding them first.
+
+    A ``bfloat16``/``float16`` product of same-dtype operands is exactly representable in float32
+    (8+8 and 11+11 mantissa bits both fit in 24) and cuBLAS already accumulates in float32, so
+    asking the GEMM for a float32 output is numerically equivalent to casting both operands up and
+    running a float32 GEMM -- but it keeps the low-precision tensor cores. Casting *after* the
+    matmul, as this used to do, is precision-inert: the logits are already rounded to the input
+    dtype by then, and JSD is a difference of nearly-equal distributions where that rounding
+    cancels catastrophically in ``dx``.
+
+    Falls back to the historical cast-after-matmul on configurations without ``out_dtype`` support.
+    The accurate alternative there needs a float32 copy of the head (V x H x 4 bytes, 2.1 GB at
+    llama shapes), which is not worth an OOM on platforms that previously worked.
+
+    Deliberately also falls back when the operands disagree, which is the autocast case: with a
+    float32 head under ``autocast(bfloat16)`` it is autocast that harmonizes the operands, and
+    ``aten::mm.dtype`` is not on its promotion list, so passing ``out_dtype`` there raises rather
+    than helping. Recovering that case means either rounding the float32 head ourselves once per
+    chunk -- duplicating autocast's weight cache -- or refusing the downcast and paying a float32
+    GEMM, which overrides the low-precision matmul the user asked autocast for. Both are trade-offs
+    rather than the free win this function makes, so they are left to a follow-up.
+    """
+    if (
+        _MM_SUPPORTS_OUT_DTYPE
+        and inputs.device.type == "cuda"
+        and torch.cuda.get_device_capability(inputs.device)[0] >= 8
+        and inputs.dtype in (torch.float16, torch.bfloat16)
+        and weight.dtype == inputs.dtype
+    ):
+        return torch.mm(inputs, weight.t(), out_dtype=torch.float32)
+    return (inputs @ weight.t()).to(torch.float32)
+
+
 # Chunk-memory constants: the chunked forward below sizes its BT chunks so
 # the transient per-chunk logits buffers live in a budget proportional to
 # BT x H. The original budget constant was 1 (budget == BT x H), which is a
@@ -143,10 +178,11 @@ def fused_linear_jsd_forward(
         teacher_input_chunk = teacher_input[start_idx:end_idx]
 
         # 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)
+        # Everything from the logits to the final JSD loss runs in FP32. The float32 output is
+        # requested from the projection GEMM itself rather than cast afterwards -- a cast after a
+        # low-precision matmul cannot recover bits the matmul already threw away.
+        student_logits_chunk = _project_to_float32(student_input_chunk, student_weight)
+        teacher_logits_chunk = _project_to_float32(teacher_input_chunk, teacher_weight)
 
         # log-softmax with temperature
         student_logits_chunk = student_logits_chunk / temperature
@@ -211,7 +247,7 @@ def fused_linear_jsd_forward(
             else:
                 grad_logits_t = student_logits_chunk.t()
                 if (
-                    _ADDMM_SUPPORTS_OUT_DTYPE
+                    _MM_SUPPORTS_OUT_DTYPE
                     and grad_weight.device.type == "cuda"
                     and torch.cuda.get_device_capability(grad_weight.device)[0] >= 8
                     and grad_weight.dtype == torch.float32

test/transformers/test_fused_linear_jsd.py

@@ -1,6 +1,7 @@
 import pytest
 import torch
 
+from packaging.version import Version
 from test.transformers.test_jsd import JSD as TorchJSD
 from test.utils import assert_verbose_allclose
 from test.utils import set_seed
@@ -483,3 +484,87 @@ def test_amp(autocast_dtype, atol, rtol):
         atol=atol,
         rtol=rtol,
     )
+
+
+def _float32_projection_reference(student_input, student_weight, teacher_input, teacher_weight, beta, temperature):
+    """Reference whose projection is *not* rounded to the input dtype.
+
+    ``TorchLMHeadJSD`` above runs its ``nn.Linear`` in the input dtype and only then casts to
+    float32, so it rounds the logits exactly the way the kernel used to. That makes it unable to
+    detect logit rounding. This oracle casts the operands first, which is exact for bfloat16 and
+    float16 operands, so the only difference from the kernel is where the rounding happens.
+    """
+    student_logits = student_input.float() @ student_weight.float().t()
+    teacher_logits = teacher_input.float() @ teacher_weight.float().t()
+    student_prob = torch.log_softmax(student_logits / temperature, dim=-1)
+    teacher_prob = torch.log_softmax(teacher_logits / temperature, dim=-1)
+    return TorchJSD(beta=beta, dtype=torch.float32)(student_prob, teacher_prob)
+
+
+# Logit spread chosen per beta so the reference gradient stays well inside the storage dtype's
+# range. The generalized-JSD mixture sits near a stationary point at beta=0.5: a wider spread there
+# drives |grad_input| to ~4e-6, below float16's smallest normal (6.1e-5), so output quantization
+# rather than the projection would dominate the error. A wider spread at beta=0.5 also makes the
+# mixture ``lerp(exp(log_q), exp(log_p), beta)`` underflow to zero and the loss go NaN -- in the
+# TorchJSD reference above as much as in the kernel, so it is not something this test can pin on
+# the projection.
+_PROJECTION_LOGIT_STD = {0.0: 30.0, 0.5: 10.0, 1.0: 30.0}
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or Version(torch.__version__.split("+")[0]) < Version("2.8.0"),
+    reason="the float32 projection needs torch>=2.8 out_dtype support on CUDA sm_80+",
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("beta", [0.0, 0.5, 1.0])
+def test_logits_projection_is_not_rounded_to_input_dtype(dtype, beta):
+    """Test that the fused projection keeps float32 logits instead of rounding them first.
+
+    Regression test for the projection matmul running in the input dtype: its output was rounded to
+    8 (bfloat16) or 11 (float16) mantissa bits before the float32 cast could preserve anything. The
+    JSD scalar hides this because it is dominated by large terms, but ``dx`` is a difference of
+    nearly-equal quantities, so the rounding cancels catastrophically. Before the fix the gradients
+    were off by 1.3-23%; after it they land within 0.7%, which is where a float32 projection sits.
+
+    The existing ``test_correctness`` cases cannot catch this: they draw both operands from
+    ``torch.rand``, which yields all-positive correlated inputs whose logits are tightly clustered.
+    Scaling a ``randn`` weight to a realistic logit spread is what exposes it.
+    """
+    if dtype is torch.float16 and beta == 0.5:
+        pytest.skip(
+            "at beta=0.5 the reference gradient peaks at ~4e-6, below float16's smallest normal, so "
+            "float16 quantization of the output buffer dominates and the projection cannot be isolated"
+        )
+
+    BT, H, V = 256, 1024, 32000
+    temperature = 1.0
+    logit_std = _PROJECTION_LOGIT_STD[beta]
+    torch.manual_seed(0)
+
+    student_input = torch.randn(BT, H, device=device, dtype=dtype)
+    teacher_input = torch.randn(BT, H, device=device, dtype=dtype)
+    student_weight = (torch.randn(V, H, device=device) * (logit_std / H**0.5)).to(dtype)
+    teacher_weight = (torch.randn(V, H, device=device) * (logit_std / H**0.5)).to(dtype)
+
+    liger_input = student_input.detach().clone().requires_grad_(True)
+    liger_weight = student_weight.detach().clone().requires_grad_(True)
+    LigerFusedLinearJSD(jsd_beta=beta, temperature=temperature)(
+        liger_input, liger_weight, teacher_input, teacher_weight, None
+    ).backward()
+
+    ref_input = student_input.detach().clone().requires_grad_(True)
+    ref_weight = student_weight.detach().clone().requires_grad_(True)
+    _float32_projection_reference(ref_input, ref_weight, teacher_input, teacher_weight, beta, temperature).backward()
+
+    for name, actual, expected in (
+        ("grad_input", liger_input.grad, ref_input.grad),
+        ("grad_weight", liger_weight.grad, ref_weight.grad),
+    ):
+        assert torch.isfinite(actual).all(), f"{name} has non-finite values"
+        relative_error = (
+            (actual.float() - expected.float()).abs().max() / expected.float().abs().max().clamp_min(1e-12)
+        ).item()
+        assert relative_error < 0.01, (
+            f"{name} deviates {relative_error:.2%} from the float32-projection reference, "
+            f"which means the projection rounded its logits to {dtype}"
+        )