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 (it does not support 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: 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

Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

…arakeet-tdt

ebezzam · 2026-05-19T06:46:36Z

    "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 1},
    "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 1},
    "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "ep-support"},
+    "tdt-loss": {"repo_id": "eustlb/tdt-loss", "revision": "v1"},


Related comment: https://github.com/huggingface/transformers/pull/44171/changes#r3094638013

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.

github-actions · 2026-09-29T08:31:06Z

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

run-slow: parakeet

github-actions · 2026-09-29T08:50:08Z

CI recap

Dashboard: View test results in Grafana
Latest run: 36543082143:1
Result: success | Jobs: 16 | Tests: 182,426 | Failures: 0 | Duration: 15h 49m

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

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

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 and others added 10 commits April 17, 2026 17:34

Merge branch 'main' into parakeet-tdt

9ab08d1

Add TDT to mapping after merge.

fd9f8b1

Fix lasr generate test.

136f676

Output attention mask if labels provided for computing loss.

833d289

Apply suggestion from @ArthurZucker

a1c62a1

Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

Improve ParakeetTDTDecoderCache definition and usage.

8683570

Remove tuple parsing.

1d4b0f4

processor refactor

a418eca

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

5d0c631

…arakeet-tdt

Update conversion.

5c603c1

ebezzam marked this pull request as draft May 19, 2026 02:33

ebezzam and others added 2 commits May 19, 2026 15:42

Merge branch 'main' into tdt_loss_kernel

09ba99c

Modular after merge.

e743b2d

ebezzam commented May 19, 2026

View reviewed changes

Don't allow all kernels.

8d09cb6

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

Open

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

Navigation Menu

Add TDT loss kernel - #46048

Add TDT loss kernel#46048
ebezzam wants to merge 89 commits into
huggingface:mainfrom
ebezzam:tdt_loss_kernel

ebezzam commented May 19, 2026 •

edited by Deep-unlearning

Loading

Uh oh!

ebezzam May 19, 2026

Uh oh!

ebezzam May 19, 2026

Uh oh!

HuggingFaceDocBuilderDev commented May 19, 2026

Uh oh!

vasqu Sep 28, 2026

Uh oh!

Deep-unlearning Sep 29, 2026

Uh oh!

vasqu Sep 28, 2026

Uh oh!

Deep-unlearning Sep 29, 2026

Uh oh!

github-actions Bot commented Sep 29, 2026

Uh oh!

github-actions Bot commented Sep 29, 2026

Uh oh!

vasqu Sep 30, 2026

Uh oh!

vasqu Sep 30, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

ebezzam commented May 19, 2026 • edited by Deep-unlearning 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!

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!

github-actions Bot commented Sep 29, 2026

Uh oh!

github-actions Bot commented Sep 29, 2026

CI recap

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

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 Deep-unlearning

Loading