Add TDT loss kernel

src/transformers/integrations/hub_kernels.py

@@ -709,6 +709,13 @@ def _build_kernel_mapping() -> dict:
                     ),
                 },
             },
+            "tdt_loss": {
+                "cuda": {
+                    Mode.TRAINING | Mode.TORCH_COMPILE: LayerRepository(
+                        repo_id="kernels-community/tdt-loss", layer_name="TDTLoss", version=1
+                    ),
+                },
+            },
         }
 
         return _KERNEL_MAPPING

src/transformers/loss/loss_tdt.py

@@ -14,12 +14,14 @@
 
 import torch
 
+from ..integrations import use_kernel_forward_from_hub
 from ..utils import logging
 
 
 logger = logging.get_logger(__name__)
 
 
+@use_kernel_forward_from_hub("tdt_loss")
 def tdt_loss(
     token_logits: torch.Tensor,
     duration_logits: torch.Tensor,

src/transformers/models/parakeet/modeling_parakeet.py

@@ -29,6 +29,7 @@
 from ...activations import ACT2FN
 from ...generation import CompileConfig, GenerationMixin, GenerationMode
 from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
+from ...loss.loss_tdt import tdt_loss
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, CausalLMOutput
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
@@ -1051,6 +1052,7 @@ def __init__(self, config: ParakeetTDTConfig):
     Parakeet Encoder with a TDT (Token Duration Transducer) head.
     """
 )
+@use_kernelized_func(tdt_loss)
 class ParakeetForTDT(ParakeetTDTGenerationMixin, ParakeetForRNNT):
     config: ParakeetTDTConfig
 

src/transformers/models/parakeet/modular_parakeet.py

@@ -23,6 +23,8 @@
 from ... import initialization as init
 from ...activations import ACT2FN
 from ...generation import CompileConfig, GenerationMixin, GenerationMode
+from ...integrations import use_kernelized_func
+from ...loss.loss_tdt import tdt_loss
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, CausalLMOutput
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
@@ -900,6 +902,7 @@ def __init__(self, config: ParakeetTDTConfig):
     Parakeet Encoder with a TDT (Token Duration Transducer) head.
     """
 )
+@use_kernelized_func(tdt_loss)
 class ParakeetForTDT(ParakeetTDTGenerationMixin, ParakeetForRNNT):
     config: ParakeetTDTConfig
 

tests/models/parakeet/test_modeling_parakeet.py

@@ -19,7 +19,16 @@
 from pathlib import Path
 
 from transformers import is_datasets_available, is_torch_available
-from transformers.testing_utils import cleanup, require_torch, require_torchaudio, slow, torch_device
+from transformers.testing_utils import (
+    cleanup,
+    preserve_module_forwards,
+    require_kernels,
+    require_torch,
+    require_torch_gpu,
+    require_torchaudio,
+    slow,
+    torch_device,
+)
 
 from ...test_configuration_common import ConfigTester
 from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
@@ -608,6 +617,41 @@ def create_and_check_model(self, config, inputs_dict):
             (self.batch_size, self.output_seq_length, self.encoder_model_tester.hidden_size),
         )
 
+    def prepare_config_and_inputs_for_loss(self):
+        config, input_features, attention_mask = self.prepare_config_and_inputs()
+        labels = ids_tensor([self.batch_size, 6], self.blank_token_id)  # no blank in the labels
+        labels[labels == self.pad_token_id] = self.pad_token_id + 1  # nor padding
+        blank = torch.full((self.batch_size, 1), self.blank_token_id, device=labels.device)
+        inputs_dict = {
+            "input_features": input_features,
+            "attention_mask": attention_mask,
+            "decoder_input_ids": torch.cat([blank, labels], dim=1),
+            "labels": labels,
+        }
+        return config, inputs_dict
+
+    def create_and_check_use_kernels(self, config, inputs_dict):
+        """`use_kernels=True` swaps `tdt_loss` for the `kernels-community/tdt-loss` kernel, with the same loss."""
+        model = ParakeetForTDT(config).to(torch_device).eval()
+        with torch.no_grad():
+            expected = model(**inputs_dict).loss
+
+        with preserve_module_forwards(model):
+            model.set_use_kernels(True)
+            self.parent.assertIn("forward", vars(tdt_loss), "`tdt_loss` was not kernelized")
+            with torch.no_grad():
+                loss = model(**inputs_dict).loss
+            torch.testing.assert_close(loss, expected, rtol=1e-4, atol=1e-4)
+        self.parent.assertNotIn("forward", vars(tdt_loss), "`tdt_loss` was not restored")
+
+        # Training mode also uses the kernel, which has a backward
+        with preserve_module_forwards(model):
+            model.train()
+            model.set_use_kernels(True)
+            self.parent.assertIn("forward", vars(tdt_loss), "`tdt_loss` was not kernelized in training mode")
+            model(**inputs_dict).loss.backward()
+            self.parent.assertTrue(any(p.grad is not None for p in model.parameters()), "No gradients after backward")
+
     def prepare_config_and_inputs_for_common(self):
         config, input_features, attention_mask = self.prepare_config_and_inputs()
         decoder_input_ids = ids_tensor([self.batch_size, 1], self.vocab_size)
@@ -650,6 +694,12 @@ def test_model(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs_for_common()
         self.model_tester.create_and_check_model(*config_and_inputs)
 
+    @require_torch_gpu
+    @require_kernels
+    def test_use_kernels(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs_for_loss()
+        self.model_tester.create_and_check_use_kernels(*config_and_inputs)
+
     @unittest.skip(reason="ParakeetForTDT does not use inputs_embeds")
     def test_model_get_set_embeddings(self):
         pass