Add TDT loss kernel by ebezzam · Pull Request #46048 · huggingface/transformers

ebezzam · 2026-05-19T02:32:47Z

What does this PR do?

Use the kernels-community/tdt-loss CUDA kernel for the TDT loss of ParakeetForTDT, for faster training.

The kernel is swapped in with the usual kernels exchange pattern (as for rotary_pos_emb):

tdt_loss is decorated with use_kernel_forward_from_hub("tdt_loss") and attached to ParakeetForTDT with use_kernelized_func
the tdt_loss mapping points to the TDTLoss layer of kernels-community/tdt-loss (version 1), for training and inference, including torch.compile
so it is used with ParakeetForTDT.from_pretrained(..., use_kernels=True); the PyTorch implementation stays the default

Speed (loss fwd+bwd, batch 4, 200 frames, 40 labels, vocab 8193): 1259 ms → 9.5 ms on A100, 993 ms → 5.0 ms on H200. Fine-tuning parakeet-tdt-0.6b-v3 on LibriSpeech gives the same loss curve as the PyTorch loss, with a full training step going from 1882 ms to 213 ms on A100 and peak memory from 33.6 to 21.8 GiB.

Tests: ParakeetForTDTModelTest.test_use_kernels checks that use_kernels=True swaps tdt_loss for the kernel (in inference and training) with the same loss. The kernel's correctness (against the PyTorch loss, a naive loop version and float64) is tested in kernels-community.

Merge order: this needs kernels-community/tdt-loss to be published first (huggingface/kernels-community#882), otherwise use_kernels=True can't load it. Tested on GPU with the staged build of the kernel (kernels-staging/tdt-loss@pr-882).

cc @eustlb @vasqu

🤖 Generated with Claude Code

Implement Token-and-Duration Transducer (TDT) decoding for Parakeet models, extending the existing CTC-only support. This adds ParakeetForTDT with greedy TDT decoding in generate(), per-token timestamp generation, and full integration with AutoModelForTDT, processors, and ASR pipeline.

- Use -100 label padding for training (HF convention) - Fix timestamp recording in inner blank-seeking loop - Add max_symbols_per_step guard matching NeMo - Clean up decoding loop - Add TDT training example to docs - Use setUpClass for TDT integration tests

…arakeet-tdt

ebezzam · 2026-05-19T06:50:14Z

+        Verify that ParakeetForTDT loss matches NeMo's TDT loss (sigma=0) for both
+        the CUDA kernel and the pure PyTorch implementation.
        reproducer: https://gist.github.com/883ea42bf7d8ce2af42f3055627476a7


Should we also test for sigma != 0?

HuggingFaceDocBuilderDev · 2026-05-19T07:01:27Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

# Conflicts: # src/transformers/integrations/hub_kernels.py # tests/models/parakeet/test_modeling_parakeet.py

- Point the hub kernel mapping to kernels-community/tdt-loss (version 1). - Dispatch to the kernel for CUDA inputs, for all reductions and sigma values. - Test the kernel against the NeMo fixtures and the PyTorch implementation (including sigma != 0), and run the integration loss test with both. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Train mode updates the BatchNorm running statistics, which changed the eval loss of the second backend. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ntegration test loops - Any error while loading the kernel (e.g. the Hub is unreachable) now falls back to the PyTorch implementation instead of failing the loss. - Remove a duplicated forward/backward in the integration test and free the outputs inside each backend's subtest. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Decorate `tdt_loss` with `use_kernel_forward_from_hub` and attach it to ParakeetForTDT with `use_kernelized_func`, so that `use_kernels=True` swaps in the `TDTLoss` layer of kernels-community/tdt-loss (as for rotary), instead of loading the kernel automatically.

vasqu · 2026-09-28T17:08:02Z

                },
            },
+            "tdt_loss": {
+                "cuda": LayerRepository(repo_id="kernels-community/tdt-loss", layer_name="TDTLoss", version=1),


iirc your kernel did not support compile right? lets give proper flags

Right, it supports backward but not torch.compile (can_torch_compile = False on the layer). Done in 3f02948: the mapping now registers it for Mode.TRAINING and Mode.INFERENCE only.

vasqu · 2026-09-28T17:09:16Z

Imo no need to test it on this side as the kernel should test it for itself no? Meaning kernels community ensures that the loss is correct 🤔

Makes sense, the kernel tests in kernels-community already check it against a PyTorch reference, a naive loop version and float64. Done in 3f02948: I removed the loss-correctness tests here and restored the integration test, keeping only test_use_kernels, which checks that use_kernels=True actually swaps tdt_loss for the kernel (inference and training) with the same loss.

Also tested against your staged build (kernels-staging/tdt-loss@pr-882) on an L4: the kernels-community tests pass (43), the Parakeet tests pass with use_kernels=True (including the slow integration test and the full test file), and the speed matches my earlier numbers (loss fwd+bwd 1285 ms → 53 ms at Parakeet size).

… the kernel exchange test - Register the kernel for training and inference only (it does not support torch.compile). - The kernel's correctness is tested in kernels-community: keep only the test that use_kernels=True swaps tdt_loss for it, and restore the integration test.

vasqu · 2026-09-30T15:59:57Z

+@require_torch
+@require_torch_gpu
+@require_kernels
+class TDTLossKernelTest(unittest.TestCase):


lets make this a test under the model instead of a separate class pls

Done in 5bd9cac: test_use_kernels is now a test of ParakeetForTDTModelTest.

vasqu · 2026-09-30T16:00:37Z

                    ),
                },
            },
+            "tdt_loss": {


We might have to drop the kernel for when we make use kernels true as default. best to make it toch compile friendly from the get go

Good point, done. The kernel now supports torch.compile (huggingface/kernels-community#882, 02647af): fake implementations for its ops, can_torch_compile = True on the layer, and tests with torch.compile(fullgraph=True) (fp32/bf16, forward and backward match eager) and torch.library.opcheck. Here, the mapping is now Mode.TRAINING | Mode.TORCH_COMPILE (5bd9cac), like the liger loss.

Checked on GPU with a build of the new kernel: a ParakeetForTDT kernelized with Mode.TRAINING | Mode.TORCH_COMPILE and a compiled tdt_loss match eager, and the Parakeet tests pass.

…its test to the model tests

vasqu · 2026-10-02T13:07:00Z

+def reset_parakeet_kernels():
+    """`use_kernels=True` replaces the forward of the shared kernelized functions: restore the PyTorch ones."""
+    from transformers.models.parakeet import modeling_parakeet
+
+    for fn in (tdt_loss, modeling_parakeet.apply_rotary_pos_emb):
+        vars(fn).pop("forward", None)


I think we already have preserve_module_forwards let's use that instead. This is super dependent on only certain parts being kernelized

Done in 0a4ccb2: removed the helper and wrapped the kernelized parts of test_use_kernels in preserve_module_forwards(model) (also checking that tdt_loss is restored afterwards).

vasqu · 2026-10-02T14:29:24Z


+    @require_torch_gpu
+    @require_kernels
+    def test_use_kernels(self):


lets make the split as in create_and_check_model / test_model then we can merge. Note that when kernels = True lands that we will temporarily disable training kernels as we need to make it compatible with the trainer

Done in 513ea4b: the logic is now in ParakeetForTDTModelTester.create_and_check_use_kernels (with prepare_config_and_inputs_for_loss), and test_use_kernels just calls it, as for test_model. Thanks for the heads-up on the training kernels!

vasqu · 2026-10-05T15:14:28Z

+        if not isinstance(tdt_loss, torch.nn.Module):
+            self.skipTest("Hub kernels are disabled (USE_HUB_KERNELS)")


Imo should not be needed because we will change the logic in any case so the function is always implicitly a module after we default to kernels true

Done in 8c0d977.

@eustlb

## What does this PR do? Goal of this PR is to add a [kernel](https://huggingface.co/eustlb/tdt-loss/tree/v1) that has been drafted by @eustlb for the recently merged [Parakeet TDT](huggingface/transformers#44171). 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](https://huggingface.co/eustlb/tdt-loss/tree/v1/build/torch211-cxx11-cu128-x86_64-linux/tdt_loss)? 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](https://github.com/huggingface/kernels-community/blob/main/CONTRIBUTING.md#how-to-add-a-new-kernel) Perhaps we can see how other kernels [used in Transformers](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/hub_kernels.py) were integrated. As the kernel is for a function, maybe [rotary_pos_emb](https://github.com/huggingface/transformers/blob/0df9b7fcaab447c75543598e6d959065c2296a24/src/transformers/integrations/hub_kernels.py#L221) is a good example? [Here](https://github.com/huggingface/kernels-community/tree/main/rotary) are the files in this repo. Here is a draft PR for adding/using this kernel in Transformers: huggingface/transformers#46048 cc @vasqu --------- Co-authored-by: Deep-unlearning <steven@huggingface.co> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Deep-unlearning · 2026-10-10T11:55:29Z

@vasqu the kernel is merged and live on the Hub (kernels-community/tdt-loss, v1) 🎉 I re-ran the Parakeet tests on an L4 against the published kernel and all pass. I also restricted test_use_kernels to the tdt_loss mapping so it doesn't depend on other kernels' builds, and fixed the Hub link. Could you take another look and trigger run-slow: parakeet when you get a chance? Thanks!

vasqu

close to being ready 🤗 im not sure whether we want to merge this or wait until after training kernels and use_kernels=True becomes a thing because it needs to be deactivated either way

vasqu · 2026-10-12T13:23:34Z

+    With `use_kernels=True`, CUDA models use the [`kernels-community/tdt-loss`](https://huggingface.co/kernels/kernels-community/tdt-loss)
+    kernel instead, which computes the same loss with fused CUDA kernels.


hmm not sure whether it makes sense to keep this, we need some general kernels docs in the future that just lists all possible kernels

vasqu · 2026-10-12T13:24:31Z

+        # Only kernelize `tdt_loss`, so that the test does not depend on the other kernels of the model
+        model.kernel_config = KernelConfig(
+            kernel_mapping={"tdt_loss": get_kernel_mapping_transformers()["tdt_loss"]}, inherit_mapping=False
+        )


might be overkill no? it should be small enough diffs to not matter

…pping restriction

Deep-unlearning · 2026-10-13T18:03:37Z

Thanks @vasqu! Removed the docstring note and the KernelConfig restriction in the test (3d23b40). On timing, happy either way: since it's opt-in, merging now lets people use the kernel already, and it's a one-line mapping change to deactivate it when use_kernels=True lands. But if you'd rather wait, that's fine too 🤗

github-actions · 2026-10-13T18:04:58Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: parakeet

github-actions · 2026-10-13T18:13:58Z

CI recap

Dashboard: View test results in Grafana
Latest run: 38443671700:1
Result: success | Jobs: 15 | Tests: 193,911 | Failures: 0 | Duration: 6h 33m

vasqu · 2026-10-14T12:46:02Z

Then I'd rather wait for now if that's ok. Trying to keep new kernels to a minimum until everything is resolved

vasqu · 2026-10-14T12:46:30Z

But feel free to ping me and follow #47773 for updates

Hainan Xu and others added 30 commits February 20, 2026 09:45

parakeet tdt intergration

fa7d6e0

Add expected outputs for TDT, small fixes.

fa36657

Separate CTC and TDT generate outputs.

05e2e34

Work with auto device, better init,

bb5ff33

Test timestamps and expose token duration.

9ec79b0

Add reproducer link.

33f128e

revert: restore lasr generated files to original state

b33002f

warn: torchaudio rnnt_loss does not train duration head

48b39dd

Relax timestamp test, and test nits.

e9f23ab

feat: TDT training

e2b97aa

chore: for cuda detection and run without patching

6b9fc73

Equivalent timestamp processing as Nemo, and various nits/cleanup.

6c879bc

Merge branch 'parakeet-tdt' of github.com:lmaksym/transformers into p…

149e17f

…arakeet-tdt

Simplify durations config.

36bfa63

Update training examples.

2df0ccc

chore: enable parralelism

388c6d3

chore: performance optimization

08b2b55

fix: formatting

0c4e05a

Doc and testing nits

1ddd804

Use active mask from current step, and nits.

f512670

Better pre-allocate.

07d8e35

TDT has separate pad token and blank token.

fab050a

Merge branch 'main' into parakeet-tdt

c438565

Regenerate lasr.

86d980c

Merge branch 'parakeet-tdt' of github.com:lmaksym/transformers into p…

895c4a0

…arakeet-tdt

Style checks and nits

ab21380

Nits, put back ctc loss test

d0141d5

More standard model output.

f7529d4

ebezzam commented May 19, 2026

View reviewed changes

ebezzam mentioned this pull request May 19, 2026

tdt-loss: add TDT loss kernel huggingface/kernels-community#882

Merged

4 tasks

Deep-unlearning and others added 6 commits September 28, 2026 14:36

Merge remote-tracking branch 'origin/main' into tdt-loss-kernel

5176486

# Conflicts: # src/transformers/integrations/hub_kernels.py # tests/models/parakeet/test_modeling_parakeet.py

Run the TDT integration loss checks before the train-mode backward

12bbc03

Train mode updates the BatchNorm running statistics, which changed the eval loss of the second backend. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Skip the kernel checks when Hub kernels are disabled

d6e739c

vasqu reviewed Sep 28, 2026

View reviewed changes

vasqu reviewed Sep 30, 2026

View reviewed changes

Address review: register the tdt_loss kernel for torch.compile, move …

5bd9cac

…its test to the model tests

vasqu reviewed Oct 2, 2026

View reviewed changes

Address review: use preserve_module_forwards in test_use_kernels

0a4ccb2

vasqu reviewed Oct 2, 2026

View reviewed changes

Address review: split test_use_kernels into the model tester

513ea4b

vasqu reviewed Oct 5, 2026

View reviewed changes

Address review: drop the kernels-disabled skip in test_use_kernels

8c0d977

Only kernelize tdt_loss in test_use_kernels, fix the kernel Hub link

3321659

vasqu reviewed Oct 12, 2026

View reviewed changes

Address review: drop the kernel note in the docstring and the test ma…

3d23b40

…pping restriction

		if not isinstance(tdt_loss, torch.nn.Module):
		self.skipTest("Hub kernels are disabled (USE_HUB_KERNELS)")

		With `use_kernels=True`, CUDA models use the [`kernels-community/tdt-loss`](https://huggingface.co/kernels/kernels-community/tdt-loss)
		kernel instead, which computes the same loss with fused CUDA kernels.

Conversation

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

Uh oh!

What does this PR do?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented May 19, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Deep-unlearning commented Oct 10, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Deep-unlearning commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026

CI recap

Uh oh!

vasqu commented Oct 14, 2026

Uh oh!

vasqu commented Oct 14, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ebezzam commented May 19, 2026 •

edited by github-actions Bot

Loading