fix(kl_div, tvd): make torch.compile work (NameError, then fp64 loop-carried scalar) (#1509)

## Summary

On current PyTorch, `torch.compile` of `LigerKLDIVLoss` and
`LigerTVDLoss` fails for every
reduction mode. There are two stacked causes; the second only shows once
the first is fixed.

**1. `NameError` from the `reduction` default.** Both kernels declare
`reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN`, a module global.
Inductor re-emits a
user-defined Triton kernel's source in its own module and brings along
the globals the body
uses, but not names that only appear in a default, so evaluating the
default fails:

```
InductorError: NameError: name '_REDUCTION_MODE_BATCHMEAN' is not defined
```

The callers always pass `reduction`, so it is now a required argument
with no default (a forgotten
argument raises `TypeError` instead of silently picking a reduction).

**2. fp64 loop-carried scalar (kl_div).** Continues the sweep offered in
#1350, same root cause:
a non-`constexpr` Python float is fp32 under eager Triton but fp64 under
Inductor. In
`_kldiv_kernel_forward`, `eps` feeds `tl.log(tl.maximum(y_true, eps))`
inside the reduction
loop, so with fix 1 alone `log_target=False` fails like #1362 did:

```
AssertionError('Loop-carried variable loss_sum has initial type fp32 but is re-assigned to fp64 in loop!')
```

`eps` (kl_div) and `scale` (tvd) are now cast with `.to(tl.float32)`
before the loop, as in
`rms_norm.py` and #1350/#1362. That is a no-op on the eager path.

The `_ascend` mirrors get the same signature change (review feedback);
the fp32 casts are not applied
there, matching the scope of #1350/#1362. I have no Ascend hardware, so
those kernels are not run.

## Testing Done

H100 80GB HBM3, torch 2.13.0+cu130, Triton 3.7.1. GPU tests ran on
`main` at `492bfeb`; the branch is rebased onto `b297821`, which does
not touch these files (`make checkstyle` re-run after the rebase).

| | `torch.compile` of kl_div / tvd |
|---|---|
| `main` (`492bfeb`) | every case fails: `NameError` above |
| fix 1 only | tvd compiles; kl_div with `log_target=False` fails with
the fp64 loop-carried error above |
| fix 1 + fix 2 (this PR) | compiles in every reduction / `log_target`
case |

- New `test_torch_compile_matches_eager` in `test_kl_div.py` (8 cases)
and `test_tvd.py` (4 cases).
The unreduced output matches eager bit for bit (`torch.equal`); reduced
losses use
  `assert_close`.
- Why not bit-identical for reduced kl_div: I compared the kernel's
per-row output (before the
wrapper's `.sum()`) between eager and compiled, and it is bit-identical
in all 6
`reduction × log_target` cases. Compiling only the `.sum()` / division
on the eager rows
reproduces the compiled loss exactly, so the remaining difference
(1.2e-7 for `batchmean`,
7.6e-6 for `sum`) comes from Inductor's own reduction order, not from
the kernel.
- `test_kl_div.py` + `test_tvd.py`: **340 passed**. Re-run on the
follow-up commit `f34c737`
  (required `reduction`, same GPU model and versions): 340 passed.
- `make checkstyle`: passes.
- `make test`: not clean, so not checked. I ran it on the same box with
the PyTorch / transformers
versions above, which are newer than CI.
`test_grpo_loss.py::test_grpo_loss_with_vllm_is_ratio_reduced[2-128-1000-0.04-luspo-sequence]`
and
`test_fused_linear_jsd.py::test_low_probability_vocab_block_is_finite`
failed. After the
`test_grpo_loss` failure, every later test in the same process errored
at setup (including
`test_kl_div.py`, which passes 340/340 when run on its own), and the run
hit my 40-minute cap
at 84%. Both failures are in files this PR does not touch. I don't have
a `main` baseline from
  that box to show they are pre-existing.

- Hardware Type: H100 80GB HBM3
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

I used an AI assistant (Claude) to help investigate and write this
change; I reviewed and ran everything above.

@vaibhavjindal

---------

Co-authored-by: Vaibhav Jindal <[email protected]>
Co-authored-by: Arup De <[email protected]>

src/liger_kernel/ops/backends/_ascend/ops/kl_div.py

@@ -37,8 +37,8 @@ def _kldiv_kernel_forward(
     eps,
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
+    reduction: tl.constexpr,  # no default: always pass explicitly
     log_target: tl.constexpr = False,
-    reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN,
 ):
     pid = tl.program_id(0)
     num_progs = tl.num_programs(0)
@@ -89,8 +89,8 @@ def _kldiv_kernel_backward(
     n_cols,
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
+    reduction: tl.constexpr,  # no default: always pass explicitly
     log_target: tl.constexpr = False,
-    reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN,
 ):
     pid = tl.program_id(0)
     num_progs = tl.num_programs(0)

src/liger_kernel/ops/backends/_ascend/ops/tvd.py

@@ -30,7 +30,7 @@ def _tv_distance_kernel(
     total_rows: tl.constexpr,  # BT
     BLOCK_SIZE: tl.constexpr,
     HAS_LABEL: tl.constexpr,
-    reduction: tl.constexpr = "batchmean",
+    reduction: tl.constexpr,  # no default: always pass explicitly
 ):
     thread_id = tl.program_id(0)
     num_threads = tl.num_programs(0)

src/liger_kernel/ops/kl_div.py

@@ -54,8 +54,10 @@ def _kldiv_kernel_forward(
     n_cols,  # int, number of columns in the input tensor
     eps,
     BLOCK_SIZE: tl.constexpr,
+    # no default: torch.compile re-emits this kernel's source in a module without the
+    # _REDUCTION_MODE_* globals, so a default naming one raises NameError. Always pass it.
+    reduction: tl.constexpr,
     log_target: tl.constexpr = False,
-    reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN,
 ):
     pid = tl.program_id(0).to(tl.int64)
     y_ptr += pid * y_stride
@@ -64,6 +66,12 @@ def _kldiv_kernel_forward(
 
     base_offsets = tl.arange(0, BLOCK_SIZE)
 
+    # Pin the scalar to fp32. Eager Triton specializes a non-constexpr Python float to
+    # fp32, but Inductor specializes it to fp64; mixing an fp64 scalar into the loop body
+    # promotes the whole expression -- including tl.log -- to float64, and re-assigns the
+    # fp32 loop-carried `loss_sum` to fp64. This cast is a no-op on the eager path.
+    eps = eps.to(tl.float32)
+
     loss_sum = 0.0
     for i in range(0, n_cols, BLOCK_SIZE):
         offsets = i + base_offsets

src/liger_kernel/ops/tvd.py

@@ -52,7 +52,9 @@ def _tv_distance_kernel(
     scale,  # pre-computed reduction scale for gradients (fused into kernel)
     BLOCK_SIZE: tl.constexpr,
     HAS_LABEL: tl.constexpr,
-    reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN,
+    # no default: torch.compile re-emits this kernel's source in a module without the
+    # _REDUCTION_MODE_* globals, so a default naming one raises NameError. Always pass it.
+    reduction: tl.constexpr,
 ):
     pid = tl.program_id(0).to(tl.int64)
     p_ptr += pid * p_stride
@@ -63,6 +65,12 @@ def _tv_distance_kernel(
 
     base_offsets = tl.arange(0, BLOCK_SIZE)
 
+    # Pin the scalar to fp32. Eager Triton specializes a non-constexpr Python float to
+    # fp32, but Inductor specializes it to fp64; mixing an fp64 scalar into the loop body
+    # promotes the gradient expression to float64 on every iteration and makes the
+    # compiled path disagree with eager. This cast is a no-op on the eager path.
+    scale = scale.to(tl.float32)
+
     if HAS_LABEL:
         label = tl.load(label_ptr)
         if label == ignore_index:

test/transformers/test_kl_div.py

@@ -102,3 +102,32 @@ def test_correctness_not_last(B, T, V, log_target, reduction, dtype, atol, rtol)
         log_target,
         is_last_layer=False,
     )
+
+
+@pytest.mark.parametrize("log_target", [True, False])
+@pytest.mark.parametrize("reduction", ["batchmean", "sum", "mean", "none"])
+def test_torch_compile_matches_eager(log_target, reduction):
+    # inductor re-emits the kernel source; a signature default that names a module global must
+    # still resolve there, and eps must stay fp32 (inductor passes Python floats as fp64). The
+    # unreduced loss matches eager bit for bit; a reduced loss may differ in the last bits
+    # because inductor fuses the wrapper's sum/division
+    torch._dynamo.reset()
+    torch.manual_seed(0)
+    liger_kldiv = LigerKLDIVLoss(reduction=reduction, log_target=log_target)
+    compiled = torch.compile(liger_kldiv)
+    input = torch.randn(64, 1271, device=device, dtype=torch.float32).log_softmax(dim=-1)
+    with torch.no_grad():
+        target = torch.randn(64, 1271, device=device).softmax(dim=-1)
+        if log_target:
+            target = target.log()
+    x1 = input.detach().clone().requires_grad_(True)
+    x2 = input.detach().clone().requires_grad_(True)
+    out1 = liger_kldiv(x1, target)
+    out2 = compiled(x2, target)
+    if reduction == "none":
+        assert torch.equal(out1, out2)
+        return
+    torch.testing.assert_close(out2, out1)
+    out1.backward()
+    out2.backward()
+    torch.testing.assert_close(x2.grad, x1.grad)

test/transformers/test_tvd.py

@@ -237,3 +237,25 @@ def test_correctness_with_ignore_index(B, T, V, reduction, dtype, atol, rtol, ig
     liger_tvd = LigerTVDLoss(reduction=reduction, ignore_index=ignore_index)
     torch_tvd = TorchTVDLoss(reduction=reduction, ignore_index=ignore_index)
     _test_correctness_with_ignore_index_once(liger_tvd, torch_tvd, ignore_index, B, T, V, dtype, atol, rtol, reduction)
+
+
+@pytest.mark.parametrize("reduction", ["batchmean", "sum", "mean", "none"])
+def test_torch_compile_matches_eager(reduction):
+    # inductor re-emits the kernel source; a signature default that names a module global must
+    # still resolve there, and the compiled call must match eager bit for bit
+    torch._dynamo.reset()
+    torch.manual_seed(0)
+    device = infer_device()
+    liger_tvd = LigerTVDLoss(reduction=reduction)
+    compiled = torch.compile(liger_tvd)
+    p = torch.randn(64, 1271, device=device).softmax(dim=-1)
+    q = torch.randn(64, 1271, device=device).softmax(dim=-1)
+    x1 = p.detach().clone().requires_grad_(True)
+    x2 = p.detach().clone().requires_grad_(True)
+    out1 = liger_tvd(x1, q)
+    out2 = compiled(x2, q)
+    assert torch.equal(out1, out2)
+    if reduction != "none":
+        out1.backward()
+        out2.backward()
+        assert torch.equal(x1.grad, x2.grad)