diff --git a/tdt-loss/CARD.md b/tdt-loss/CARD.md
new file mode 100644
index 00000000..1b383245
--- /dev/null
+++ b/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 %}
diff --git a/tdt-loss/README.md b/tdt-loss/README.md
new file mode 100644
index 00000000..52260115
--- /dev/null
+++ b/tdt-loss/README.md
@@ -0,0 +1,47 @@
+---
+tags:
+- kernel
+---
+This CUDA extension implements TDT (Token-and-Duration Transducer) loss, originally from [eustlb/tdt-loss](https://huggingface.co/eustlb/tdt-loss).
+
+## Usage
+
+```python
+# /// script
+# dependencies = [
+#   "torch",
+#   "kernels"
+# ]
+# ///
+import torch
+from kernels import get_kernel
+
+tdt_loss_module = get_kernel("kernels-community/tdt-loss")
+
+# Example usage
+batch_size = 2
+max_input_len = 100       # T: number of encoder frames
+max_target_len = 20       # U: includes blank column (= num_labels + 1)
+vocab_size = 256
+durations = [0, 1, 2, 3, 4]  # possible duration values
+num_durations = len(durations)
+blank_id = 0
+device = torch.device("cuda")
+
+token_logits = torch.randn(batch_size, max_input_len, max_target_len, vocab_size, device=device, requires_grad=True)
+duration_logits = torch.randn(batch_size, max_input_len, max_target_len, num_durations, device=device, requires_grad=True)
+targets = torch.randint(1, vocab_size, (batch_size, max_target_len - 1), device=device)
+source_lengths = torch.full((batch_size,), max_input_len, device=device, dtype=torch.int32)
+target_lengths = torch.full((batch_size,), max_target_len - 1, device=device, dtype=torch.int32)
+
+loss = tdt_loss_module.tdt_loss(
+    token_logits, duration_logits, targets,
+    source_lengths, target_lengths, durations,
+    blank_id, sigma=0.05, reduction="mean",
+)
+loss.backward()
+```
+
+## Benchmarks
+
+See the [benchmark script](https://gist.github.com/eustlb/42bb54363e82ad6c0f6fc8e24f604c18).
diff --git a/tdt-loss/build.toml b/tdt-loss/build.toml
new file mode 100644
index 00000000..bb3444af
--- /dev/null
+++ b/tdt-loss/build.toml
@@ -0,0 +1,30 @@
+[general]
+name = "tdt-loss"
+license = "Apache-2.0"
+version = 1
+backends = ["cuda"]
+
+[general.hub]
+repo-id = "kernels-community/tdt-loss"
+
+[torch]
+src = [
+    "torch-ext/torch_binding.cpp",
+    "torch-ext/torch_binding.h",
+]
+
+[kernel.tdt_loss]
+backend = "cuda"
+cuda-capabilities = [
+    "8.0",
+    "8.9",
+    "9.0",
+    "10.0",
+    "12.0",
+]
+depends = ["torch"]
+include = ["."]
+src = [
+    # TODO: add CUDA source files here, e.g.:
+    # "csrc/tdt_loss.cu",
+]
diff --git a/tdt-loss/flake.nix b/tdt-loss/flake.nix
new file mode 100644
index 00000000..ef8a89a2
--- /dev/null
+++ b/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 = ./.;
+    };
+}
diff --git a/tdt-loss/tests/test_tdt_loss.py b/tdt-loss/tests/test_tdt_loss.py
new file mode 100644
index 00000000..236982e6
--- /dev/null
+++ b/tdt-loss/tests/test_tdt_loss.py
@@ -0,0 +1,221 @@
+"""Tests for TDT loss kernel, validated against a naive loop-based reference."""
+
+import pytest
+import torch
+
+if not torch.cuda.is_available():
+    pytest.skip("CUDA not available", allow_module_level=True)
+
+try:
+    from tdt_loss import tdt_loss
+except ImportError as e:
+    pytest.skip(f"tdt_loss not available: {e}", allow_module_level=True)
+
+
+# ---------------------------------------------------------------------------
+# Naive loop-based reference (no vectorisation, easy to verify by hand)
+# ---------------------------------------------------------------------------
+
+def _tdt_loss_naive(
+    token_logits,       # (B, T, U, V)
+    duration_logits,    # (B, T, U, D)
+    targets,            # (B, U-1)
+    source_lengths,     # (B,)
+    target_lengths,     # (B,)
+    durations,          # list[int]
+    blank_id,
+    sigma=0.0,
+):
+    """Single-sample, triple-nested-loop reference. Returns per-sample losses."""
+    B, max_T, max_U, V = token_logits.shape
+    D = len(durations)
+    device = token_logits.device
+
+    token_logits = token_logits.float()
+    duration_logits = duration_logits.float()
+
+    tok_lp = torch.log_softmax(token_logits, dim=-1) - sigma
+    dur_lp = torch.log_softmax(duration_logits, dim=-1)
+
+    losses = []
+    for b in range(B):
+        T_b = int(source_lengths[b].item())
+        U_b = int(target_lengths[b].item()) + 1
+
+        alpha = torch.full((T_b, U_b), float("-inf"), device=device, dtype=torch.float64)
+        alpha[0, 0] = 0.0
+
+        for t in range(T_b):
+            for u in range(U_b):
+                if t == 0 and u == 0:
+                    continue
+                val = torch.tensor(float("-inf"), device=device, dtype=torch.float64)
+                for i, dur in enumerate(durations):
+                    t_src = t - dur
+                    if t_src < 0:
+                        continue
+                    # Blank arc (dur > 0): (t_src, u) -> (t, u)
+                    if dur > 0:
+                        arc = alpha[t_src, u] + tok_lp[b, t_src, u, blank_id] + dur_lp[b, t_src, u, i]
+                        val = torch.logaddexp(val, arc.to(torch.float64))
+                    # Label arc (any dur): (t_src, u-1) -> (t, u)
+                    if u > 0:
+                        label = int(targets[b, u - 1].item())
+                        arc = alpha[t_src, u - 1] + tok_lp[b, t_src, u - 1, label] + dur_lp[b, t_src, u - 1, i]
+                        val = torch.logaddexp(val, arc.to(torch.float64))
+                alpha[t, u] = val
+
+        # Terminal: blank arcs with dur > 0 that exit the lattice
+        ll = torch.tensor(float("-inf"), device=device, dtype=torch.float64)
+        for i, dur in enumerate(durations):
+            if dur == 0:
+                continue
+            t_src = T_b - dur
+            if t_src < 0:
+                continue
+            arc = alpha[t_src, U_b - 1] + tok_lp[b, t_src, U_b - 1, blank_id] + dur_lp[b, t_src, U_b - 1, i]
+            ll = torch.logaddexp(ll, arc.to(torch.float64))
+
+        losses.append(-ll)
+
+    return torch.stack(losses)
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+@pytest.mark.parametrize("sigma", [0.0, 0.05])
+@pytest.mark.parametrize("reduction", ["mean", "sum", "none"])
+def test_tdt_loss_forward(sigma, reduction):
+    """Forward loss matches the naive loop-based reference."""
+    torch.manual_seed(42)
+    device = torch.device("cuda")
+
+    B, T, U_labels, V = 4, 30, 8, 32
+    durations = [0, 1, 2, 3, 4]
+    D = len(durations)
+    max_U = U_labels + 1
+    blank_id = 0
+
+    token_logits = torch.randn(B, T, max_U, V, device=device)
+    duration_logits = torch.randn(B, T, max_U, D, device=device)
+    targets = torch.randint(1, V, (B, U_labels), device=device)
+    source_lengths = torch.tensor([T, T - 2, T - 5, T], device=device, dtype=torch.int32)
+    target_lengths = torch.tensor(
+        [U_labels, U_labels - 1, U_labels - 2, U_labels], device=device, dtype=torch.int32
+    )
+
+    # Naive reference (float64 for accuracy)
+    ref_per_sample = _tdt_loss_naive(
+        token_logits, duration_logits, targets,
+        source_lengths, target_lengths, durations, blank_id,
+        sigma=sigma,
+    ).float()
+
+    # Kernel under test
+    kernel_out = tdt_loss(
+        token_logits, duration_logits, targets,
+        source_lengths, target_lengths, durations,
+        blank_id, sigma=sigma, reduction=reduction,
+    )
+
+    if reduction == "none":
+        ref = ref_per_sample
+    elif reduction == "mean":
+        ref = (ref_per_sample / target_lengths.float()).mean()
+    else:
+        ref = ref_per_sample.sum()
+
+    torch.testing.assert_close(kernel_out.cpu(), ref.cpu(), atol=1e-4, rtol=1e-4)
+
+
+@pytest.mark.parametrize("sigma", [0.0, 0.05])
+def test_tdt_loss_backward(sigma):
+    """Backward gradients are finite and consistent between two identical calls."""
+    torch.manual_seed(123)
+    device = torch.device("cuda")
+
+    B, T, U_labels, V = 2, 20, 5, 16
+    durations = [0, 1, 2, 3]
+    D = len(durations)
+    max_U = U_labels + 1
+    blank_id = 0
+
+    tok_data = torch.randn(B, T, max_U, V, device=device, dtype=torch.float32)
+    dur_data = torch.randn(B, T, max_U, D, device=device, dtype=torch.float32)
+    targets = torch.randint(1, V, (B, U_labels), device=device)
+    source_lengths = torch.tensor([T, T - 3], device=device, dtype=torch.int32)
+    target_lengths = torch.tensor([U_labels, U_labels - 1], device=device, dtype=torch.int32)
+
+    # Run 1
+    tok1 = tok_data.clone().requires_grad_(True)
+    dur1 = dur_data.clone().requires_grad_(True)
+    loss1 = tdt_loss(
+        tok1, dur1, targets, source_lengths, target_lengths,
+        durations, blank_id, sigma=sigma, reduction="mean",
+    )
+    loss1.backward()
+
+    # Run 2 (should produce identical gradients)
+    tok2 = tok_data.clone().requires_grad_(True)
+    dur2 = dur_data.clone().requires_grad_(True)
+    loss2 = tdt_loss(
+        tok2, dur2, targets, source_lengths, target_lengths,
+        durations, blank_id, sigma=sigma, reduction="mean",
+    )
+    loss2.backward()
+
+    # Losses should match
+    torch.testing.assert_close(loss1, loss2)
+
+    # Gradients should be identical and finite
+    assert torch.isfinite(tok1.grad).all(), "token_logits gradients contain non-finite values"
+    assert torch.isfinite(dur1.grad).all(), "duration_logits gradients contain non-finite values"
+    torch.testing.assert_close(tok1.grad, tok2.grad)
+    torch.testing.assert_close(dur1.grad, dur2.grad)
+
+
+def test_tdt_loss_batch_size_one():
+    """Sanity check with batch size 1."""
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+
+    B, T, U_labels, V = 1, 15, 4, 10
+    durations = [0, 1, 2]
+    D = len(durations)
+    max_U = U_labels + 1
+    blank_id = 0
+
+    token_logits = torch.randn(B, T, max_U, V, device=device)
+    duration_logits = torch.randn(B, T, max_U, D, device=device)
+    targets = torch.randint(1, V, (B, U_labels), device=device)
+    source_lengths = torch.tensor([T], device=device, dtype=torch.int32)
+    target_lengths = torch.tensor([U_labels], device=device, dtype=torch.int32)
+
+    ref = _tdt_loss_naive(
+        token_logits, duration_logits, targets,
+        source_lengths, target_lengths, durations, blank_id,
+    ).float()
+
+    out = tdt_loss(
+        token_logits, duration_logits, targets,
+        source_lengths, target_lengths, durations,
+        blank_id, reduction="none",
+    )
+    torch.testing.assert_close(out.cpu(), ref.cpu(), atol=1e-4, rtol=1e-4)
+
+
+def test_tdt_loss_invalid_reduction():
+    """Should raise on invalid reduction."""
+    device = torch.device("cuda")
+    B, T, U, V, D = 1, 5, 3, 4, 2
+    with pytest.raises(ValueError, match="Invalid reduction"):
+        tdt_loss(
+            torch.randn(B, T, U, V, device=device),
+            torch.randn(B, T, U, D, device=device),
+            torch.randint(1, V, (B, U - 1), device=device),
+            torch.tensor([T], device=device, dtype=torch.int32),
+            torch.tensor([U - 1], device=device, dtype=torch.int32),
+            [0, 1], 0, reduction="invalid",
+        )
diff --git a/tdt-loss/torch-ext/tdt_loss/__init__.py b/tdt-loss/torch-ext/tdt_loss/__init__.py
new file mode 100644
index 00000000..8649c7bb
--- /dev/null
+++ b/tdt-loss/torch-ext/tdt_loss/__init__.py
@@ -0,0 +1,149 @@
+"""TDT (Token-and-Duration Transducer) loss CUDA kernel."""
+
+from typing import List, Union
+
+import torch
+
+from ._ops import ops
+
+
+class TDTLoss(torch.autograd.Function):
+    """Custom autograd function for TDT loss."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        token_logits: torch.Tensor,
+        duration_logits: torch.Tensor,
+        targets: torch.Tensor,
+        source_lengths: torch.Tensor,
+        target_lengths: torch.Tensor,
+        durations: torch.Tensor,
+        blank_id: int,
+        sigma: float = 0.0,
+    ) -> torch.Tensor:
+        B, max_T, max_U, V = token_logits.shape
+        D = duration_logits.shape[3]
+        device = token_logits.device
+
+        token_logits = token_logits.contiguous().float()
+        duration_logits = duration_logits.contiguous().float()
+        targets = targets.contiguous().int()
+        source_lengths = source_lengths.contiguous().int()
+        target_lengths = target_lengths.contiguous().int()
+        durations = durations.contiguous().int()
+
+        blank_lp = torch.empty(B, max_T, max_U, device=device, dtype=torch.float32)
+        label_lp = torch.empty(B, max_T, max_U, device=device, dtype=torch.float32)
+        dur_lp = torch.empty(B, max_T, max_U, D, device=device, dtype=torch.float32)
+        alphas = torch.full((B, max_T, max_U), -1e30, device=device, dtype=torch.float32)
+        log_ll = torch.empty(B, device=device, dtype=torch.float32)
+
+        ops.tdt_logprobs_fwd(
+            token_logits, duration_logits, targets,
+            source_lengths, target_lengths, blank_id, sigma,
+            blank_lp, label_lp, dur_lp,
+        )
+        ops.tdt_loss_fwd(
+            blank_lp, label_lp, dur_lp,
+            source_lengths, target_lengths, durations,
+            alphas, log_ll,
+        )
+
+        ctx.save_for_backward(
+            token_logits, duration_logits, targets,
+            source_lengths, target_lengths, durations,
+            blank_lp, label_lp, dur_lp, alphas, log_ll,
+        )
+        ctx.blank_id = blank_id
+        return -log_ll
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (token_logits, duration_logits, targets,
+         source_lengths, target_lengths, durations,
+         blank_lp, label_lp, dur_lp, alphas, log_ll) = ctx.saved_tensors
+        blank_id = ctx.blank_id
+
+        B, max_T, max_U, V = token_logits.shape
+        D = duration_logits.shape[3]
+        device = token_logits.device
+
+        betas = torch.full((B, max_T, max_U), -1e30, device=device, dtype=torch.float32)
+        ll_bwd = torch.empty(B, device=device, dtype=torch.float32)
+        ops.tdt_loss_bwd(
+            blank_lp, label_lp, dur_lp,
+            source_lengths, target_lengths, durations,
+            betas, ll_bwd,
+        )
+
+        grad_blank = torch.zeros(B, max_T, max_U, device=device, dtype=torch.float32)
+        grad_label = torch.zeros(B, max_T, max_U, device=device, dtype=torch.float32)
+        grad_dur = torch.zeros(B, max_T, max_U, D, device=device, dtype=torch.float32)
+        ops.tdt_loss_grad(
+            alphas, betas, blank_lp, label_lp, dur_lp, log_ll,
+            source_lengths, target_lengths, durations,
+            grad_blank, grad_label, grad_dur,
+        )
+
+        grad_blank = grad_blank * grad_output.unsqueeze(-1).unsqueeze(-1)
+        grad_label = grad_label * grad_output.unsqueeze(-1).unsqueeze(-1)
+        grad_dur = grad_dur * grad_output.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+
+        grad_token_logits = torch.zeros_like(token_logits)
+        grad_duration_logits = torch.zeros_like(duration_logits)
+        ops.tdt_logprobs_bwd(
+            token_logits, duration_logits, targets,
+            source_lengths, target_lengths, blank_id,
+            grad_blank, grad_label, grad_dur,
+            grad_token_logits, grad_duration_logits,
+        )
+
+        return grad_token_logits, grad_duration_logits, None, None, None, None, None, None
+
+
+def tdt_loss(
+    token_logits: torch.Tensor,
+    duration_logits: torch.Tensor,
+    targets: torch.Tensor,
+    source_lengths: torch.Tensor,
+    target_lengths: torch.Tensor,
+    durations: Union[List[int], torch.Tensor],
+    blank_id: int,
+    sigma: float = 0.0,
+    reduction: str = "mean",
+) -> torch.Tensor:
+    """Compute TDT (Token-and-Duration Transducer) loss using CUDA kernels.
+
+    Args:
+        token_logits: Token logits of shape (batch, T, U+1, vocab_size+1).
+        duration_logits: Duration logits of shape (batch, T, U+1, num_durations).
+        targets: Target labels of shape (batch, U).
+        source_lengths: Encoder output lengths of shape (batch,).
+        target_lengths: Target lengths of shape (batch,).
+        durations: List or 1-D tensor of duration values (e.g. [0, 1, 2, 3, 4]).
+        blank_id: Blank token id.
+        sigma: Logit undernormalization constant (see TDT paper). Defaults to 0.0.
+        reduction: Loss reduction method: "mean", "sum", or "none".
+
+    Returns:
+        Scalar loss tensor (or per-example losses if reduction="none").
+    """
+    if reduction not in ("mean", "sum", "none"):
+        raise ValueError(f'Invalid reduction mode "{reduction}". Expected one of "mean", "sum", or "none".')
+
+    if isinstance(durations, (list, tuple)):
+        durations = torch.tensor(durations, dtype=torch.int32, device=token_logits.device)
+    else:
+        durations = durations.to(device=token_logits.device, dtype=torch.int32).contiguous()
+
+    per_sample_loss = TDTLoss.apply(
+        token_logits, duration_logits, targets,
+        source_lengths, target_lengths, durations,
+        blank_id, sigma,
+    )
+    if reduction == "mean":
+        return (per_sample_loss / target_lengths.float()).mean()
+    elif reduction == "sum":
+        return per_sample_loss.sum()
+    return per_sample_loss
diff --git a/tdt-loss/torch-ext/torch_binding.cpp b/tdt-loss/torch-ext/torch_binding.cpp
new file mode 100644
index 00000000..5b7ca2ee
--- /dev/null
+++ b/tdt-loss/torch-ext/torch_binding.cpp
@@ -0,0 +1,42 @@
+#include <torch/library.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def(
+      "tdt_logprobs_fwd(Tensor token_logits, Tensor duration_logits, "
+      "Tensor targets, Tensor source_lengths, Tensor target_lengths, "
+      "int blank_id, float sigma, "
+      "Tensor! blank_lp, Tensor! label_lp, Tensor! dur_lp) -> ()");
+  ops.impl("tdt_logprobs_fwd", torch::kCUDA, &tdt_logprobs_fwd);
+
+  ops.def(
+      "tdt_loss_fwd(Tensor blank_lp, Tensor label_lp, Tensor dur_lp, "
+      "Tensor source_lengths, Tensor target_lengths, Tensor durations, "
+      "Tensor! alphas, Tensor! log_ll) -> ()");
+  ops.impl("tdt_loss_fwd", torch::kCUDA, &tdt_loss_fwd);
+
+  ops.def(
+      "tdt_loss_bwd(Tensor blank_lp, Tensor label_lp, Tensor dur_lp, "
+      "Tensor source_lengths, Tensor target_lengths, Tensor durations, "
+      "Tensor! betas, Tensor! ll_bwd) -> ()");
+  ops.impl("tdt_loss_bwd", torch::kCUDA, &tdt_loss_bwd);
+
+  ops.def(
+      "tdt_loss_grad(Tensor alphas, Tensor betas, Tensor blank_lp, "
+      "Tensor label_lp, Tensor dur_lp, Tensor log_ll, "
+      "Tensor source_lengths, Tensor target_lengths, Tensor durations, "
+      "Tensor! grad_blank, Tensor! grad_label, Tensor! grad_dur) -> ()");
+  ops.impl("tdt_loss_grad", torch::kCUDA, &tdt_loss_grad);
+
+  ops.def(
+      "tdt_logprobs_bwd(Tensor token_logits, Tensor duration_logits, "
+      "Tensor targets, Tensor source_lengths, Tensor target_lengths, "
+      "int blank_id, "
+      "Tensor grad_blank, Tensor grad_label, Tensor grad_dur, "
+      "Tensor! grad_token_logits, Tensor! grad_duration_logits) -> ()");
+  ops.impl("tdt_logprobs_bwd", torch::kCUDA, &tdt_logprobs_bwd);
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
diff --git a/tdt-loss/torch-ext/torch_binding.h b/tdt-loss/torch-ext/torch_binding.h
new file mode 100644
index 00000000..1e5b68f6
--- /dev/null
+++ b/tdt-loss/torch-ext/torch_binding.h
@@ -0,0 +1,53 @@
+#pragma once
+
+#include <torch/torch.h>
+
+// Forward: compute log-probabilities for blank, label, and duration
+void tdt_logprobs_fwd(torch::Tensor const &token_logits,
+                      torch::Tensor const &duration_logits,
+                      torch::Tensor const &targets,
+                      torch::Tensor const &source_lengths,
+                      torch::Tensor const &target_lengths, int64_t blank_id,
+                      double sigma, torch::Tensor &blank_lp,
+                      torch::Tensor &label_lp, torch::Tensor &dur_lp);
+
+// Forward pass of TDT loss (alpha computation)
+void tdt_loss_fwd(torch::Tensor const &blank_lp,
+                  torch::Tensor const &label_lp,
+                  torch::Tensor const &dur_lp,
+                  torch::Tensor const &source_lengths,
+                  torch::Tensor const &target_lengths,
+                  torch::Tensor const &durations, torch::Tensor &alphas,
+                  torch::Tensor &log_ll);
+
+// Backward pass of TDT loss (beta computation)
+void tdt_loss_bwd(torch::Tensor const &blank_lp,
+                  torch::Tensor const &label_lp,
+                  torch::Tensor const &dur_lp,
+                  torch::Tensor const &source_lengths,
+                  torch::Tensor const &target_lengths,
+                  torch::Tensor const &durations, torch::Tensor &betas,
+                  torch::Tensor &ll_bwd);
+
+// Gradient computation
+void tdt_loss_grad(torch::Tensor const &alphas, torch::Tensor const &betas,
+                   torch::Tensor const &blank_lp,
+                   torch::Tensor const &label_lp,
+                   torch::Tensor const &dur_lp, torch::Tensor const &log_ll,
+                   torch::Tensor const &source_lengths,
+                   torch::Tensor const &target_lengths,
+                   torch::Tensor const &durations,
+                   torch::Tensor &grad_blank, torch::Tensor &grad_label,
+                   torch::Tensor &grad_dur);
+
+// Backward: propagate gradients through log-probabilities to logits
+void tdt_logprobs_bwd(torch::Tensor const &token_logits,
+                      torch::Tensor const &duration_logits,
+                      torch::Tensor const &targets,
+                      torch::Tensor const &source_lengths,
+                      torch::Tensor const &target_lengths, int64_t blank_id,
+                      torch::Tensor const &grad_blank,
+                      torch::Tensor const &grad_label,
+                      torch::Tensor const &grad_dur,
+                      torch::Tensor &grad_token_logits,
+                      torch::Tensor &grad_duration_logits);