tdt-loss: add TDT loss kernel

.github/kernel-maintainers.json

@@ -51,6 +51,7 @@
   "scattermoe": ["U08Q42K8MGW"],
   "sgl-flash-attn3": [],
   "sonic-moe": ["U08Q42K8MGW", "U054KJK89C2"],
+  "tdt-loss": ["U07U53UJECT"],
   "tinygrad-rms": [],
   "topk": ["U0586KAH2E9"],
   "trimul-gpumode": [],

scripts/check_kernel_freshness.py

@@ -48,6 +48,7 @@
     "scattermoe": "https://github.com/shawntan/scattermoe",
     "sgl-flash-attn3": "https://github.com/sgl-project/sgl-flash-attn",
     "sonic-moe": "https://github.com/Dao-AILab/sonic-moe",
+    "tdt-loss": "",
     "tinygrad-rms": "https://github.com/tinygrad/tinygrad",
     "trimul-gpumode": "https://github.com/davidberard98/gpumode-trimul",
     "triton-kernels": "https://github.com/triton-lang/triton.git",

tdt-loss/CARD.md

@@ -0,0 +1,42 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}---
+
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+{% if functions %}
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+kernel_module = get_kernel("{{ repo_id }}")
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
+
+{{ functions[0] }}(...)
+```
+{% else %}
+
+Usage example not available.
+{% endif %}
+
+## Available functions
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+
+## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }}`.
+{% else %}
+
+Benchmarking script is not available.
+{% endif %}

tdt-loss/benchmarks/benchmark.py

@@ -0,0 +1,81 @@
+import sys
+from pathlib import Path
+
+import torch
+
+from kernels.benchmark import Benchmark
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from tests.reference import tdt_loss_reference  # noqa: E402
+
+DURATIONS = [0, 1, 2, 3, 4]
+
+
+class TDTLossBenchmark(Benchmark):
+    seed: int = 42
+
+    def _make_inputs(self, B, T, U, V, dtype):
+        D = len(DURATIONS)
+        # Slices of a joint output, as produced by ParakeetForTDT.
+        self.logits = torch.randn(B, T, U + 1, V + D, device=self.device, dtype=dtype, requires_grad=True)
+        self.targets = torch.randint(1, V, (B, U), device=self.device, dtype=torch.int32)
+        self.logit_lengths = torch.full((B,), T, device=self.device, dtype=torch.int32)
+        self.target_lengths = torch.full((B,), U, device=self.device, dtype=torch.int32)
+        self.V = V
+
+    def _run(self):
+        self.logits.grad = None
+        losses = self.kernel.tdt_loss(
+            self.logits[..., : self.V],
+            self.logits[..., self.V :],
+            self.targets,
+            self.logit_lengths,
+            self.target_lengths,
+            0,
+            DURATIONS,
+            reduction="none",
+        )
+        losses.sum().backward()
+        self.out = losses.detach()
+
+    def _reference(self):
+        with torch.no_grad():
+            return tdt_loss_reference(
+                self.logits[..., : self.V],
+                self.logits[..., self.V :],
+                self.targets,
+                self.logit_lengths,
+                self.target_lengths,
+                0,
+                DURATIONS,
+            )
+
+    # Small vocabulary, short utterances.
+    def setup(self):
+        self._make_inputs(B=8, T=100, U=20, V=1025, dtype=torch.float32)
+
+    def benchmark_base(self):
+        self._run()
+
+    def verify_base(self) -> torch.Tensor:
+        return self._reference()
+
+    # Parakeet TDT v3 sizes (8192 tokens + blank), float32.
+    def setup_large(self):
+        self._make_inputs(B=4, T=200, U=40, V=8193, dtype=torch.float32)
+
+    def benchmark_large(self):
+        self._run()
+
+    def verify_large(self) -> torch.Tensor:
+        return self._reference()
+
+    # Same sizes with bfloat16 logits.
+    def setup_large_bf16(self):
+        self._make_inputs(B=4, T=200, U=40, V=8193, dtype=torch.bfloat16)
+
+    def benchmark_large_bf16(self):
+        self._run()
+
+    def verify_large_bf16(self) -> torch.Tensor:
+        return self._reference()

tdt-loss/build.toml

@@ -0,0 +1,28 @@
+[general]
+name = "tdt-loss"
+license = "Apache-2.0"
+version = 1
+edition = 5
+backends = ["cuda"]
+
+[general.hub]
+repo-id = "kernels-community/tdt-loss"
+
+[torch]
+src = [
+    "torch-ext/torch_binding.cpp",
+    "torch-ext/torch_binding.h",
+]
+
+[torch.stable-abi]
+cuda = "2.10"
+
+[kernel.tdt_loss]
+backend = "cuda"
+depends = ["torch"]
+src = [
+    "tdt_loss/common.cuh",
+    "tdt_loss/lattice.cu",
+    "tdt_loss/logprobs.cu",
+    "tdt_loss/stable.h",
+]

tdt-loss/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for tdt-loss kernel";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernels";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}