Fix TVD gradient at p == q

[truong-v]

Summary

_tv_distance_kernel computes the gradient as tl.where(p > q, 0.5 * scale, -0.5 * scale), a two-way split of a three-way sign: p == q falls into the p < q branch and gets -0.5 * scale where the derivative of 0.5 * |p - q| is 0, so the loss is unaffected but the gradient is wrong on exact ties. Ties are reachable with default settings — self-distillation (student and teacher identical at step 0) makes every element a tie, and padded vocabulary slots are exactly 0.0 in both distributions — and the suite's own reference, TorchTVDLoss (torch.abs(p - q) / 2.0), already returns 0 there, but the tests never draw a tie because both distributions are random.

Fixes #1373

Details

• src/liger_kernel/ops/tvd.py: wrap the existing tl.where in a zero branch for p == q. Nothing else changes — untied elements take the same path as before.
• test/transformers/test_tvd.py: new test_correctness_with_ties, which ties the trailing slots of the target (fraction 1.0 = self-distillation, 0.5 = padded vocabulary) and compares against TorchTVDLoss across all four reductions and both dtypes.

Gradient at a tie, p == q everywhere, reduction="batchmean", (8, 512) fp32 on H100:

	gradient at a tied element	wrong elements
before	-0.0625	4096 / 4096
after	0.0	0 / 4096
TorchTVDLoss reference	0.0	—

The extra comparison is one predicate and one select per element in a bandwidth-bound kernel. Forward op, (4096, 32000) bf16, do_bench, three interleaved runs on one H100 NVL: 0.2337 / 0.2338 / 0.2335 ms before, 0.2337 / 0.2339 / 0.2354 ms after — within run-to-run noise.

Testing Done

• test/transformers/test_tvd.py: 232 passed (200 pre-existing + 32 new).

• Reverting only the kernel change makes 24 of the 32 new cases fail; the 8 that still pass are reduction="none", which returns before the backward pass since the loss itself is correct.

• test_tvd.py is the only file in the suite that exercises this kernel, so nothing else in make test is affected. I ran ruff check . and ruff format --check . over the repo, but not the full suite or the convergence suite locally — boxes below reflect that.

• Hardware Type: H100 NVL

• run make test to ensure correctness

• run make checkstyle to ensure code style

• run make test-convergence to ensure convergence

Environment

• Liger-Kernel a5d795efd2c1436549e70118ef519134e9c27833 (main) + this branch, editable install
• GPU: NVIDIA H100 NVL
• torch 2.6.0+cu124, triton 3.2.0, CUDA 12.4, transformers 5.14.1, Python 3.10.20