tdt-loss: add TDT loss kernel by ebezzam · Pull Request #882 · huggingface/kernels-community

ebezzam · 2026-05-19T07:51:37Z

What does this PR do?

Goal of this PR is to add a kernel that has been drafted by @eustlb for the recently merged Parakeet TDT.

It's a very rough draft with a lot of Claude help as it's my first time working with kernels 😅

TODO

@eustlb: I guess you had a .cu file to create the .so file here? From what I understand, we will need to add it here.
Add the kernel to https://huggingface.co/kernels-community
Update the tests
Double-check with this checklist

Perhaps we can see how other kernels used in Transformers were integrated. As the kernel is for a function, maybe rotary_pos_emb is a good example? Here are the files in this repo.

Here is a draft PR for adding/using this kernel in Transformers: huggingface/transformers#46048

cc @vasqu

github-actions · 2026-05-19T07:51:56Z

Hi @ebezzam, thanks for your interest in contributing!

This project requires that pull request authors are vouched, and you are not in the list of vouched users.

This PR will be closed automatically. See https://github.com/huggingface/kernels-community/blob/main/CONTRIBUTING.md for more details.

Implement the TDT loss kernels from scratch: - logprobs.cu: log-softmax gather over the vocabulary (online softmax, one block per lattice node) and a fused gradient kernel w.r.t. the logits. - lattice.cu: alpha/beta recursions over anti-diagonals, one block per sample with threads striding over u (no U <= 1024 limit). Logits are read in float32/float16/bfloat16 with arbitrary strides on the (batch, T, U) dimensions, so slices of the joint output need no copy. Supports all NeMo reductions (mean_volume, mean_batch, mean, sum, none). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

alpha + beta - log_likelihood cancels in float32 when the log-likelihood is large (long targets). The recursions are ~2% of the runtime, so accumulate them in float64; log-probs and logits stay in float32 or lower. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Only the log-values are float64; exp/log act on differences in float32, and each node uses a two-pass logsumexp instead of chained log_add calls. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Move the implementation from __init__.py to loss.py, add layers.TDTLoss (a stateless layer with the same signature as transformers' tdt_loss) and use the transformers argument order in the functional API.

Deep-unlearning · 2026-09-28T16:26:57Z

Status / next steps

CUDA kernels, tests (43 passing on L4/A100/H200), benchmarks, README
layers.TDTLoss for the kernels exchange pattern (thanks @vasqu)
Registered in kernel-maintainers.json and check_kernel_freshness.py
Build: could someone with write access run /kernel-bot build tdt-loss? This is the first kernel-builder/nix build of this kernel (I could only test with the PyTorch JIT builder).
Review and merge, then publish kernels-community/tdt-loss v1
Then merge Add TDT loss kernel transformers#46048 (it needs the published kernel for use_kernels=True)

The "Comment with build instructions" check fails with a 403 from the bot's token. It does the same on other fork PRs (e.g. #1182).

vasqu · 2026-09-28T17:06:44Z

/kernel-bot build-and-stage tdt-loss

github-actions · 2026-09-28T17:07:03Z

Build request processed.

Command: /kernel-bot build-and-stage tdt-loss
Mode: build and stage
Target branch: pr-882
PR head SHA: 3b9907c58a98e3874b16b454a72fe8fb2a8b67c7
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

tdt-loss (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/36455858473

Hub uploads:

Hub repo: https://huggingface.co/kernels/kernels-staging/tdt-loss

vasqu · 2026-09-28T17:07:22Z

Using the staged build so you can use it via hf hub as well

Register fake implementations for the ops (they only write to preallocated outputs), build the durations tensor without the eager cache while compiling, and mark the TDTLoss layer as compile-friendly. Test with torch.compile(fullgraph=True) and torch.library.opcheck.

Deep-unlearning · 2026-10-02T10:04:00Z

@vasqu I added torch.compile support (02647af, following your comment on huggingface/transformers#46048): fake implementations for the ops, can_torch_compile = True on the TDTLoss layer, and torch.compile(fullgraph=True) + torch.library.opcheck tests (46 tests pass on an L4). Could you re-run /kernel-bot build-and-stage tdt-loss? The staged build is from before this change.

github-actions · 2026-10-02T10:04:08Z

This PR changes source code for the following kernel(s): tdt-loss.

If you have the kernel-bot build permissions, then trigger /kernel-bot build tdt-loss yourself. Otherwise, a maintainer will review the PR and will take an appropriate action. Thanks for your patience.

vasqu · 2026-10-02T13:05:09Z

/kernel-bot build-and-stage tdt-loss

github-actions · 2026-10-02T13:05:29Z

Build request processed.

Command: /kernel-bot build-and-stage tdt-loss
Mode: build and stage
Target branch: pr-882
PR head SHA: 02647aff4a7c7a6640823f9b298465844c86bb73
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

tdt-loss (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/37010709634

Hub uploads:

Hub repo: https://huggingface.co/kernels/kernels-staging/tdt-loss

Deep-unlearning · 2026-10-05T14:40:28Z

@vasqu friendly ping: all comments are addressed here and on huggingface/transformers#46048 (torch.compile support + the test split), and the new build passed. Could you approve/merge when you get a chance, so #46048 can follow? Thanks!

vasqu · 2026-10-05T15:15:23Z

  "scattermoe": ["U08Q42K8MGW"],
  "sgl-flash-attn3": [],
  "sonic-moe": ["U08Q42K8MGW", "U054KJK89C2"],
+  "tdt-loss": [],


Can you add yourself here Steven? Iirc it's the slack id

Done in ab47fc9 (U07U53UJECT).

vasqu · 2026-10-05T15:16:34Z

Not sure how much sense the readme make iirc only the card will be used to create the readmes

Makes sense, removed in ab47fc9.

vasqu · 2026-10-05T15:18:12Z

Would it be possible to port to torch ABI stable, e.g. see #1180

With this we can support more torch versions. Otherwise we will be stuck with the latest 2 which is not good and can lead to discrepancies

Done: ported to the stable ABI in febfe21 + 474a55f, following #1180 ([torch.stable-abi] cuda = "2.10", torch::stable::Tensor, STD_TORCH_CHECK, THO dispatch, stable DeviceGuard, AOTI stream, STABLE_TORCH_LIBRARY). The CUDA kernels are unchanged.

Tested on an L4 with -DTORCH_STABLE_ONLY -DTORCH_TARGET_VERSION=0x020a000000000000: all 46 tests (including torch.compile and opcheck) pass on both Torch 2.10 and 2.11, same speed. Could you re-run /kernel-bot build-and-stage tdt-loss?

Use torch::stable::Tensor, STD_TORCH_CHECK, THO dispatch, the stable DeviceGuard and the AOTI shim for the current stream, and register the ops with STABLE_TORCH_LIBRARY. The CUDA kernels are unchanged.

vasqu · 2026-10-05T16:41:29Z

/kernel-bot build-and-stage tdt-loss

github-actions · 2026-10-05T16:41:45Z

Build request processed.

Command: /kernel-bot build-and-stage tdt-loss
Mode: build and stage
Target branch: pr-882
PR head SHA: 474a55f7115008cac1ef6cd8e5ba3c8fd2197310
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

tdt-loss (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/37342798875

Deep-unlearning · 2026-10-06T09:08:21Z

Thanks @vasqu! The build failed because build.toml was missing edition = 5. Fixed in 1eefa6f, and I merged main so the checks are green again. Could you re-run /kernel-bot build-and-stage tdt-loss? Once it's merged, a release would let huggingface/transformers#46048 follow 🙏

vasqu · 2026-10-06T09:24:02Z

/kernel-bot build-and-stage tdt-loss

github-actions · 2026-10-06T09:24:18Z

Build request processed.

Command: /kernel-bot build-and-stage tdt-loss
Mode: build and stage
Target branch: pr-882
PR head SHA: 1eefa6f976267c980fb91022fd524a1d9ee6c924
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

tdt-loss (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/37442712627

Hub uploads:

Hub repo: https://huggingface.co/kernels/kernels-staging/tdt-loss

Deep-unlearning · 2026-10-09T13:03:37Z

Hi @vasqu @danieldk @drbh! The staged build on 1eefa6f passes (build + GPU tests), and all review comments are addressed. Would you be able to approve and merge when you have a moment, and trigger the release? That would unblock huggingface/transformers#46048. Thanks a lot 🙏

drbh · 2026-10-09T16:35:28Z

+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from tests.reference import tdt_loss_reference  # noqa: E402


we should update this to use the kernels library like the test, not blocking since its only for benching, but should probably be followed up in the next pr that touches this kernel

drbh

lgtm! nice addition!

drbh · 2026-10-09T16:37:10Z

/kernel-bot merge-and-upload tdt-loss

github-actions · 2026-10-09T16:37:25Z

Build request processed.

Command: /kernel-bot merge-and-upload tdt-loss
Mode: merge, build and upload
Target branch: ``
PR head SHA: 1eefa6f976267c980fb91022fd524a1d9ee6c924
Workflows: `build.yaml, build-mac.yaml, build-windows.yaml`

Merge result: Pull Request successfully merged

Dispatched (1):

tdt-loss (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/37960459170

Hub uploads:

Hub repo: https://huggingface.co/kernels/kernels-community/tdt-loss

ebezzam requested review from danieldk and drbh as code owners May 19, 2026 07:51

github-actions Bot closed this May 19, 2026

ebezzam changed the title ~~Initial commit, missing .cu file.~~ Add TDT loss kernel May 19, 2026

ebezzam mentioned this pull request May 19, 2026

Add TDT loss kernel huggingface/transformers#46048

Draft

danieldk reopened this May 19, 2026

danieldk changed the title ~~Add TDT loss kernel~~ tdt-loss: add TDT loss kernel May 19, 2026

ebezzam and others added 7 commits September 28, 2026 14:26

Initial commit, missing .cu file.

64218fb

tdt-loss: add benchmark

0bee960

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

tdt-loss: fix test references for unreachable nodes and long targets

3447945

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

tdt-loss: keep float32 transcendentals in the lattice recursions

3301b5d

Only the log-values are float64; exp/log act on differences in float32, and each node uses a two-pass logsumexp instead of chained log_add calls. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

tdt-loss: update README

da05018

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

tdt-loss: add a TDTLoss layer for the kernels exchange pattern

3b9907c

Move the implementation from __init__.py to loss.py, add layers.TDTLoss (a stateless layer with the same signature as transformers' tdt_loss) and use the transformers argument order in the functional API.

vasqu reviewed Oct 5, 2026

View reviewed changes

Deep-unlearning added 3 commits October 5, 2026 17:33

tdt-loss: add maintainer, drop README (the card is used for the Hub)

ab47fc9

tdt-loss: move to the Torch stable ABI

febfe21

Use torch::stable::Tensor, STD_TORCH_CHECK, THO dispatch, the stable DeviceGuard and the AOTI shim for the current stream, and register the ops with STABLE_TORCH_LIBRARY. The CUDA kernels are unchanged.

tdt-loss: stable ABI fixes for Torch 2.10

474a55f

Deep-unlearning added 2 commits October 6, 2026 10:52

Merge main

18a1056

tdt-loss: use build.toml edition 5, bump kernel-builder

1eefa6f

drbh reviewed Oct 9, 2026

View reviewed changes

drbh approved these changes Oct 9, 2026

View reviewed changes

github-actions Bot merged commit 8bdd234 into huggingface:main Oct 9, 2026
9 checks passed

		sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
		from tests.reference import tdt_loss_reference # noqa: E402

Conversation

ebezzam commented May 19, 2026

What does this PR do?

Uh oh!

github-actions Bot commented May 19, 2026

Uh oh!

Deep-unlearning commented Sep 28, 2026

Uh oh!

vasqu commented Sep 28, 2026

Uh oh!

github-actions Bot commented Sep 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vasqu commented Sep 28, 2026

Uh oh!

Deep-unlearning commented Oct 2, 2026

Uh oh!

github-actions Bot commented Oct 2, 2026

Uh oh!

vasqu commented Oct 2, 2026

Uh oh!

github-actions Bot commented Oct 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Deep-unlearning commented Oct 5, 2026

Uh oh!

vasqu Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

Deep-unlearning Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

Deep-unlearning Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

Deep-unlearning Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu commented Oct 5, 2026

Uh oh!

github-actions Bot commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Deep-unlearning commented Oct 6, 2026

Uh oh!

vasqu commented Oct 6, 2026

Uh oh!

github-actions Bot commented Oct 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Deep-unlearning commented Oct 9, 2026

Uh oh!

drbh Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

drbh left a comment

Choose a reason for hiding this comment

Uh oh!

drbh commented Oct 9, 2026

Uh oh!

github-actions Bot commented Oct 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

github-actions Bot commented Sep 28, 2026 •

edited

Loading

github-actions Bot commented Oct 2, 2026 •

edited

Loading

github-actions Bot commented Oct 5, 2026 •

edited

Loading

github-actions Bot commented Oct 6, 2026 •

edited

Loading

github-actions Bot commented Oct 9, 2026 •

edited

Loading