feat: add FusedLinearKLDivLoss (fused linear + KL divergence for distillation) by Yulong-Cauli · Pull Request #1423 · linkedin/Liger-Kernel

Yulong-Cauli · 2026-08-30T03:41:01Z

Summary

Add FusedLinearKLDivLoss, which fuses the vocabulary-head linear projection with
KL(Q || P) loss against an explicit target distribution, for knowledge distillation
workflows where teacher probabilities are available directly (e.g. cached / offline
soft labels), complementing LigerKLDIVLoss (element-wise, no fusion) and
LigerFusedLinearJSD (teacher side needs its own projection).

The kernel follows the chunked gradient-in-forward structure of fused_linear_jsd.py,
so the full BT x V logits tensor is never materialized, and adopts the
_CHUNK_MEM_CONST = 16 token-chunk geometry of fused_linear_cross_entropy.py.

Details

Supports reduction ("batchmean" / "mean" / "sum", same semantics as
torch.nn.KLDivLoss), ignore_index via shift_labels, softmax temperature,
and an eps clamp so that 0 * log(0) contributes 0.
Supports accum_dtype for fp32 weight-gradient accumulation across chunked
low-precision GEMMs (same semantics as LigerFusedLinearJSD /
LigerFusedLinearCrossEntropyLoss), including the torch.addmm(out_dtype=fp32)
fast path on torch >= 2.8 / sm80+ and a fallback elsewhere.
The log-softmax backward is computed in place on the saved log-probs and the
chunk_size x V fp32 temporaries are freed eagerly, so the peak during the
trailing GEMMs stays low.
reduction="none" is intentionally unsupported (gradient-in-forward requires a
scalar loss); a ValueError points users to LigerKLDIVLoss instead.
bf16 + reduction="sum" test tolerances follow the existing convention in
test_fused_linear_cross_entropy.py.
Peak-memory improvement over the unfused PyTorch baseline (full fwd+bwd,
max_memory_allocated, bf16, llama-3 head shape H=4096 / V=128256, RTX 4060
Laptop GPU): BT=2048 6537 -> 4310 MB (1.52x), BT=4096 12072 -> 6594 MB (1.83x);
at BT=1024 the benchmark harness reports 4793 -> 4244 MB. Runtime is at parity
with the baseline at BT=1024 (0.89x).

Testing Done

Hardware Type: NVIDIA GeForce RTX 4060 Laptop GPU (sm89); NVIDIA Tesla T4 (sm75)
run make test to ensure correctness (32 pre-existing failures unrelated to
this change: 18 test_mlp.py shared-memory OOR on sm89 and 14 test_grpo_loss.py
LUSPO numerics; both reproduce on main without this change)
run make checkstyle to ensure code style
run make test-convergence to ensure convergence (does not exercise the new
standalone loss)

Test logs (RTX 4060 Laptop, torch 2.11.0+cu128, triton 3.7.1)

make test: 32 failed, 4152 passed, 914 skipped, 15 xfailed in 1050.50s (0:17:30)
  test/transformers/test_fused_linear_kl_div.py: 172 passed, 0 failed, 0 skipped
  failures (all pre-existing on this GPU, reproduced on main without this change):
    18 x test/transformers/test_mlp.py
        triton.runtime.errors.OutOfResources: shared memory
        Required: 131072, Hardware limit: 101376 (sm89 has 100 KB usable smem)
    14 x test/chunked_loss/test_grpo_loss.py
        chunked GRPO LUSPO numerics (also seen by other contributors, see #1397)
make checkstyle: ruff check and ruff format --check pass on all files in this PR

Test logs (Kaggle Tesla T4, torch 2.10.0+cu128, triton 3.6.0)

test/transformers/test_fused_linear_kl_div.py: 51 passed, 25 skipped in 10.29s
(bf16 cases skipped by design: T4 is sm75 without native bf16 support)

…nsients eagerly The historical inc_factor = cdiv(V, H) formula (memory budget C=1) forces tiny 32-row chunks at LLM vocab shapes, making the chunk loop launch-bound: measured full fwd+bwd at llama-3 head shape (H=4096, V=128256, bf16, BT=1024) was 5.1x slower than the unfused baseline. Adopt fused_linear_cross_entropy's _CHUNK_MEM_CONST = 16 budget, which reaches speed parity (0.89x). Also compute the log-softmax backward in place and drop chunk_size x V fp32 temporaries (logits/softmax/log-prob) as soon as they are consumed, so the peak during the trailing GEMMs drops. Peak memory for full fwd+bwd (max_memory_allocated, RTX 4060 Laptop, bf16, llama-3 head shape): BT=2048 6537 -> 4310 MB (1.52x), BT=4096 12072 -> 6594 MB (1.83x), and at BT=1024 the fused op now also wins (4793 -> 4244 MB via the benchmark harness). Per-row loss is partition-invariant across chunk sizes; gradients only differ at the GEMM reduction-order level, same class as existing chunked fused ops.

Yulong-Cauli · 2026-09-06T14:56:28Z

Rebased onto latest main (6459426). The only conflict was one import line
in transformers/functional.py (the new LigerFusedLinearScaledCrossEntropyTPFunction
export from #1428); both exports kept in alphabetical order.

I also reviewed the new multi-DSL dispatcher architecture from #1416 and the
CuTeDSL FusedLinearJSD backend from #1419 — LigerFusedLinearKLDivLoss is
complementary to that work (KL vs JSD divergence, explicit target distribution
for cached/offline distillation). Legacy-style ops remain fully supported, so
this PR lands the Triton implementation as-is.

Follow-up plan: register the inner KL primitive via declare_op_locations
(following the jsd_loss_and_grad pattern) so a future CuTeDSL/cuTile backend
can be opted into, mirroring how fused_linear_jsd is wired.

Tested on latest main (RTX 4060 Laptop 8GB, torch 2.11.0+cu128):

test/transformers/test_fused_linear_kl_div.py: 76 passed
Sibling test_fused_linear_jsd.py + new upstream test/ops/test_kl_div.py: 83 passed, 3 skipped
ruff check / ruff format --check: clean

Could a maintainer approve the workflow run and take a look when you get a
chance? Thanks!

Adopt the same three-branch weight-gradient accumulation that fused_linear_cross_entropy uses after linkedin#1454: on CUDA SM80+ with dtype-matched low-precision params (accum_dtype=None), accumulate straight into grad_weight via addmm(out=grad_weight) instead of materializing a parameter-sized product + cast per token chunk. The fp32-accumulator out_dtype fast path and the legacy fallback are unchanged in behavior. Also expose accum_dtype in LigerFusedLinearKLDivLoss.extra_repr.

Yulong-Cauli · 2026-09-13T06:27:30Z

@arde171 @kolehma8 @BYHsu — friendly ping for a review when you get a chance.

One update since my last comment:

New commit — dW accumulation aligned with #1454. The low-precision weight-gradient accumulation now mirrors the three-branch addmm pattern #1454 adopted for FLCE: on CUDA SM80+ with dtype-matched bf16/fp16 params (accum_dtype=None), dW accumulates straight into grad_weight via torch.addmm(..., out=grad_weight) instead of materializing a parameter-sized product + cast on every token chunk. The fp32-accumulator out_dtype fast path is unchanged. Also exposed accum_dtype in extra_repr.

Validation (same three suites as before: test_fused_linear_kl_div.py, sibling test_fused_linear_jsd.py, upstream test/ops/test_kl_div.py):

RTX 4060 Laptop 8GB (SM 8.9, torch 2.11.0+cu128): 159 passed, 3 skipped — this exercises both addmm branches
Kaggle Tesla T4 (SM 7.5, torch 2.10+cu128): 134 passed, 28 skipped (bf16 cases are skipped on T4) — independent-GPU regression check
ruff clean; merges cleanly onto latest main (95b01e9, verified with git merge-tree)

As before, the PR is complementary to the JSD backends from #1419/#1420 (KL vs JSD divergence; the teacher distribution is given explicitly, e.g. cached soft labels for offline distillation), and registering the inner KL primitive via declare_op_locations (following the jsd_loss_and_grad pattern) remains the planned follow-up.

Could a maintainer approve the workflow run (checks are pending approval for fork PRs) and take a look when you get a chance? Thanks!

kolehma8 · 2026-09-13T15:03:39Z

+        chunk_n_rows = logits_chunk.shape[0]
+
+        # log-softmax with temperature
+        log_prob_chunk = torch.log_softmax(logits_chunk, dim=-1).contiguous()


Could this be fused with the main kernel (_kl_div_kernel)?

kolehma8 · 2026-09-13T15:04:07Z

+        # For anything starting from logits to the final KL loss, we do computation
+        # in FP32 to avoid losing numerical stability.
+        logits_chunk = (input_chunk @ student_weight.t()).to(torch.float32)
+        logits_chunk.div_(temperature)


Can this be moved to fused kernel?

kolehma8 · 2026-09-13T15:04:30Z

+        # shape: chunk_size x V
+        # For anything starting from logits to the final KL loss, we do computation
+        # in FP32 to avoid losing numerical stability.
+        logits_chunk = (input_chunk @ student_weight.t()).to(torch.float32)


could you do the casting to FP32 inside the kernel (e.g. registers or SMEM) to avoid HBM bloat?

kolehma8 · 2026-09-13T15:05:10Z

+        # (log_prob_chunk now holds g = dL/dlog_prob; done in place to avoid
+        # materializing extra chunk_size x V temporaries)
+        softmax_chunk = torch.softmax(logits_chunk, dim=-1)
+        del logits_chunk  # free chunk_size x V fp32 before the GEMMs below


I believe these will become unnecessary once you move some of the operations inside the fused kernel?

kolehma8 · 2026-09-13T15:07:15Z

@@ -0,0 +1,163 @@
+import torch


Can you include some performance numbers for the kernel you added compared with vanilla torch implementation?

kolehma8 · 2026-09-13T15:09:08Z

@arde171 @kolehma8 @BYHsu — friendly ping for a review when you get a chance.

One update since my last comment:

New commit — dW accumulation aligned with #1454. The low-precision weight-gradient accumulation now mirrors the three-branch addmm pattern #1454 adopted for FLCE: on CUDA SM80+ with dtype-matched bf16/fp16 params (accum_dtype=None), dW accumulates straight into grad_weight via torch.addmm(..., out=grad_weight) instead of materializing a parameter-sized product + cast on every token chunk. The fp32-accumulator out_dtype fast path is unchanged. Also exposed accum_dtype in extra_repr.

Validation (same three suites as before: test_fused_linear_kl_div.py, sibling test_fused_linear_jsd.py, upstream test/ops/test_kl_div.py):

RTX 4060 Laptop 8GB (SM 8.9, torch 2.11.0+cu128): 159 passed, 3 skipped — this exercises both addmm branches

Kaggle Tesla T4 (SM 7.5, torch 2.10+cu128): 134 passed, 28 skipped (bf16 cases are skipped on T4) — independent-GPU regression check

ruff clean; merges cleanly onto latest main (95b01e9, verified with git merge-tree)

As before, the PR is complementary to the JSD backends from #1419/#1420 (KL vs JSD divergence; the teacher distribution is given explicitly, e.g. cached soft labels for offline distillation), and registering the inner KL primitive via declare_op_locations (following the jsd_loss_and_grad pattern) remains the planned follow-up.

Could a maintainer approve the workflow run (checks are pending approval for fork PRs) and take a look when you get a chance? Thanks!

@Yulong-Cauli thank you for your contribution. I left few comments about the kernel itself and could you provide also some performance numbers? We want to ensure that all the kernels in Liger are faster (or more memory efficient) than plain torch implementations.

… KL kernel Address review feedback: the GEMM output now stays in its native precision in HBM and _kl_div_kernel upcasts to FP32 in registers, computing the temperature scaling, the log-softmax, the KL loss and the log-softmax backward in-kernel, then overwrites the logits buffer in place with dL/dlogits. This removes the fp32 logits cast, the log_softmax materialization and the torch-side softmax-recompute/backprop passes around the kernel: each token chunk now has a single native-precision transient (previously three fp32 ones) and far fewer HBM round trips. The loss is accumulated in the algebraically equal factored form sum(q*(log(max(q,eps)) - x/T)) + lse*sum(q) so it completes within the same two-pass structure.

Yulong-Cauli · 2026-09-13T17:26:55Z

@kolehma8 thanks for the review! All five points are addressed in ef19754:

The GEMM output now stays in its native precision in HBM. _kl_div_kernel upcasts to FP32 in registers on load and computes the temperature scaling, the log-softmax (two-pass max/logsumexp), the KL loss and the log-softmax backward in-kernel, overwriting the logits buffer in place with dL/dlogits. So the FP32 cast (L159), the temperature div_ (L160), the log_softmax materialization (L164) and the torch-side softmax-recompute block with its temporaries (L193) are all gone — each token chunk now has a single native-precision transient instead of three fp32 ones. The loss is accumulated in the algebraically equal factored form sum(q*(log(max(q,eps)) - x/T)) + lse*sum(q) so it fits the same two-pass structure.

Benchmark (full fwd+bwd, median of 3×30 timed iterations, incremental peak-allocated memory; methodology as in #1454):

A100-SXM4-80GB, bf16, torch 2.11.0:

BT × H × V	torch	liger	speedup	peak mem
8192 × 4096 × 128256	195.3 ms	129.2 ms	1.51x	19.0 GB → 2.0 GB (−89.3%)
4096 × 4096 × 128256	94.8 ms	64.4 ms	1.47x	9.0 GB → 1.0 GB (−88.7%)
2048 × 4096 × 128256	47.3 ms	33.1 ms	1.43x	3992 → 509 MB (−87.2%)
8192 × 2048 × 32000	33.9 ms	17.3 ms	1.96x	4843 → 532 MB (−89.0%)
4096 × 2048 × 32000	17.3 ms	9.0 ms	1.92x	2359 → 266 MB (−88.7%)

Tesla T4, fp16, torch 2.10.0 (bf16 is unsupported on SM75):

BT × H × V	torch	liger	speedup	peak mem
4096 × 2048 × 32000	121.8 ms	72.0 ms	1.69x	2358 → 625 MB (−73.5%)
2048 × 4096 × 128256	362.2 ms	354.0 ms	1.02x	3992 → 3257 MB (−18.4%)

fp32 is ~parity (0.94x–1.16x depending on shape and GPU). The fused kernel also validates on three GPUs (SM75/SM80/SM89): the full PR suite plus the sibling JSD and upstream kl_div suites pass everywhere (bf16 cases included on SM80/SM89).

Raw run logs:

A100-SXM4-80GB (SM80), torch 2.11.0+cu130, repo @ ef19754

GPU: NVIDIA A100-SXM4-80GB | torch 2.11.0+cu130
$ python -m pytest test/transformers/test_fused_linear_kl_div.py \
      test/transformers/test_fused_linear_jsd.py test/ops/test_kl_div.py -q
159 passed, 3 skipped, 37 warnings in 32.69s

$ python bench_kl_a100.py
BT=2048 H=4096 V=128256 bfloat16: torch 47.25 ms | liger 33.11 ms (1.43x) | peak 3992 -> 509 MB (-87.2%) | |loss diff| 3.67e-02
BT=4096 H=4096 V=128256 bfloat16: torch 94.80 ms | liger 64.41 ms (1.47x) | peak 8986 -> 1018 MB (-88.7%) | |loss diff| 8.41e-02
BT=8192 H=4096 V=128256 bfloat16: torch 195.27 ms | liger 129.16 ms (1.51x) | peak 18974 -> 2036 MB (-89.3%) | |loss diff| 4.42e-02
BT=4096 H=2048 V=32000 bfloat16: torch 17.27 ms | liger 8.99 ms (1.92x) | peak 2359 -> 266 MB (-88.7%) | |loss diff| 5.19e-02
BT=8192 H=2048 V=32000 bfloat16: torch 33.85 ms | liger 17.29 ms (1.96x) | peak 4843 -> 532 MB (-89.0%) | |loss diff| 3.93e-02
BT=2048 H=4096 V=128256 float32: torch 403.42 ms | liger 430.93 ms (0.94x) | peak 1972 -> 2505 MB (-27.0%) | |loss diff| 3.81e-06
BT=4096 H=2048 V=32000 float32: torch 116.39 ms | liger 113.13 ms (1.03x) | peak 1718 -> 750 MB (-56.3%) | |loss diff| 3.81e-06
DONE

Note: peak memory is the incremental max_memory_allocated above the live
baseline captured right before the timed loop; autograd's reused weight-grad
buffer stays counted in the torch baseline, while the fused op allocates its
weight-grad inside the measured window, so the fp32 rows above understate the
fused kernel's memory advantage.

Tesla T4 (SM75), torch 2.10.0+cu128, repo @ ef19754

GPU: Tesla T4 capability: (7, 5) | torch 2.10.0+cu128
$ python -m pytest test/transformers/test_fused_linear_kl_div.py \
      test/transformers/test_fused_linear_jsd.py test/ops/test_kl_div.py -v
134 passed, 28 skipped, 31 warnings in 26.04s   (bf16 cases skip on SM75)

$ python bench (same harness)
BT=2048 H=4096 V=128256 float16: torch 362.24 ms | liger 353.98 ms (1.02x) | peak 3992 -> 3257 MB (-18.4%) | |loss diff| 5.69e-03
BT=4096 H=4096 V=128256 float16: OOM — skipped
BT=8192 H=4096 V=128256 float16: OOM — skipped
BT=4096 H=2048 V=32000 float16: torch 121.76 ms | liger 72.03 ms (1.69x) | peak 2358 -> 625 MB (-73.5%) | |loss diff| 6.09e-03
BT=2048 H=4096 V=128256 float32: OOM — skipped
BT=4096 H=2048 V=32000 float32: torch 461.07 ms | liger 427.63 ms (1.08x) | peak 1718 -> 750 MB (-56.3%) | |loss diff| 3.81e-06
DONE

The OOM rows are the vanilla-torch pipeline exceeding the 16 GB card; the
whole case is skipped at that point, so the fused kernel was not measured on
those shapes on this card.

Yulong-Cauli added 2 commits September 6, 2026 22:35

feat(kl_div): add fused linear KL divergence loss

294e027

Yulong-Cauli force-pushed the feat/fused-linear-kl-div branch from 53aeedf to 2a1a964 Compare September 6, 2026 14:46

kolehma8 reviewed Sep 13, 2026

View reviewed changes

kolehma8 enabled auto-merge September 14, 2026 02:33

kolehma8 approved these changes Sep 14, 2026

View reviewed changes

kolehma8 added this pull request to the merge queue Sep 14, 2026

Merged via the queue into linkedin:main with commit f6385e9 Sep 14, 2026
1 of 3 checks passed

Navigation Menu

Uh oh!

feat: add FusedLinearKLDivLoss (fused linear + KL divergence for distillation) - #1423

feat: add FusedLinearKLDivLoss (fused linear + KL divergence for distillation)#1423
kolehma8 merged 4 commits into
linkedin:mainfrom
Yulong-Cauli:feat/fused-linear-kl-div

Yulong-Cauli commented Aug 30, 2026

Uh oh!

Yulong-Cauli commented Sep 6, 2026

Uh oh!

Yulong-Cauli commented Sep 13, 2026

Uh oh!

kolehma8 Sep 13, 2026

Uh oh!

kolehma8 Sep 13, 2026

Uh oh!

kolehma8 Sep 13, 2026

Uh oh!

kolehma8 Sep 13, 2026

Uh oh!

kolehma8 Sep 13, 2026

Uh oh!

kolehma8 commented Sep 13, 2026

Uh oh!

Yulong-Cauli commented Sep 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

Yulong-Cauli commented Aug 30, 2026

Summary

Details

Testing Done

Uh oh!

Yulong-Cauli commented Sep 6, 2026

Uh oh!

Yulong-Cauli commented Sep 13, 2026

Uh oh!

kolehma8 Sep 13, 2026

Choose a reason for hiding this comment

Uh oh!

kolehma8 Sep 13, 2026

Choose a reason for hiding this comment

Uh oh!

kolehma8 Sep 13, 2026

Choose a reason for hiding this comment

Uh oh!

kolehma8 Sep 13, 2026

Choose a reason for hiding this comment

Uh oh!

kolehma8 Sep 13, 2026

Choose a reason for hiding this comment

Uh oh!

kolehma8 commented Sep 13, 2026

Uh oh!

Yulong-Cauli commented Sep 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants