Skip to content

Add fused CE + TV distance kernel - #1384

Open
kashif wants to merge 2 commits into
linkedin:mainfrom
kashif:kashif/fused-ce-tvd
Open

Add fused CE + TV distance kernel#1384
kashif wants to merge 2 commits into
linkedin:mainfrom
kashif:kashif/fused-ce-tvd

Conversation

@kashif

@kashif kashif commented Aug 15, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds LigerFusedCETVDFunction: per-token cross-entropy and total variation distance against a teacher, computed in one streaming pass over the vocabulary.

Distillation losses that pair a hard-label term with a distributional distance currently materialize two (BT, V) float32 softmax temporaries. This fuses both terms and keeps only O(BT) state for backward, rebuilding the distributions from logits the caller already holds. ops/tvd.py is unchanged — it consumes probabilities, so it cannot avoid those temporaries.

Both terms are returned unreduced, so the caller keeps ownership of masking, per-token weighting, and normalization.

Testing Done

test/transformers/test_fused_ce_tvd.py: forward and backward parity against eager autograd in fp32 and bf16, over vocab sizes 128–131072 (multi-tile) and non-power-of-two tails, plus ignore_index, teacher-receives-no-gradient, and shape validation.

Exact p == q ties are covered in the backward pass, including self-distillation: the sign uses a zero subgradient at ties so the kernel matches torch.abs backward. Swapping it for a two-way p > q split fails both tie tests — the same bug class #1374 fixes in ops/tvd.py.

At V=151936 bf16, against a chunked and activation-checkpointed eager reference: marginal memory falls from ~0.87 MB/token to ~0 (890 MiB at 1024 tokens), and HBM traffic ~5.9x.

  • Hardware Type: AMD Radeon 890M (ROCm)
  • run make test to ensure correctness
  • run make checkstyle to ensure code style
  • run make test-convergence to ensure convergence

make test reached 2741 passed / 858 skipped before stopping on test_fused_moe.py::test_correctness, which is unrelated and simply too slow on this integrated GPU; tests ordered after it did not run. The timings above come from that same box, so the traffic ratio is the transferable number rather than the wall-clock. Please run the full suite and convergence in CI.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant