From 64218fbbbd1c36ab0c86562b87da969792efe428 Mon Sep 17 00:00:00 2001 From: Eric B Date: Tue, 19 May 2026 09:28:15 +0200 Subject: [PATCH 01/17] Initial commit, missing .cu file. --- tdt-loss/CARD.md | 42 +++++ tdt-loss/README.md | 47 +++++ tdt-loss/build.toml | 30 ++++ tdt-loss/flake.nix | 17 ++ tdt-loss/tests/test_tdt_loss.py | 221 ++++++++++++++++++++++++ tdt-loss/torch-ext/tdt_loss/__init__.py | 149 ++++++++++++++++ tdt-loss/torch-ext/torch_binding.cpp | 42 +++++ tdt-loss/torch-ext/torch_binding.h | 53 ++++++ 8 files changed, 601 insertions(+) create mode 100644 tdt-loss/CARD.md create mode 100644 tdt-loss/README.md create mode 100644 tdt-loss/build.toml create mode 100644 tdt-loss/flake.nix create mode 100644 tdt-loss/tests/test_tdt_loss.py create mode 100644 tdt-loss/torch-ext/tdt_loss/__init__.py create mode 100644 tdt-loss/torch-ext/torch_binding.cpp create mode 100644 tdt-loss/torch-ext/torch_binding.h diff --git a/tdt-loss/CARD.md b/tdt-loss/CARD.md new file mode 100644 index 000000000..1b3832451 --- /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 000000000..522601157 --- /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 000000000..bb3444af3 --- /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 000000000..ef8a89a25 --- /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 000000000..236982e6a --- /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 000000000..8649c7bbe --- /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 000000000..5b7ca2ee1 --- /dev/null +++ b/tdt-loss/torch-ext/torch_binding.cpp @@ -0,0 +1,42 @@ +#include + +#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 000000000..1e5b68f6b --- /dev/null +++ b/tdt-loss/torch-ext/torch_binding.h @@ -0,0 +1,53 @@ +#pragma once + +#include + +// 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); From 0d6fe8ddd28fbdf80ddf00a172c5c21994c9ee62 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 28 Sep 2026 14:28:04 +1000 Subject: [PATCH 02/17] tdt-loss: add CUDA kernels Implement the TDT loss kernels from scratch: - logprobs.cu: log-softmax gather over the vocabulary (online softmax, one block per lattice node) and a fused gradient kernel w.r.t. the logits. - lattice.cu: alpha/beta recursions over anti-diagonals, one block per sample with threads striding over u (no U <= 1024 limit). Logits are read in float32/float16/bfloat16 with arbitrary strides on the (batch, T, U) dimensions, so slices of the joint output need no copy. Supports all NeMo reductions (mean_volume, mean_batch, mean, sum, none). Co-Authored-By: Claude Opus 5.5 --- tdt-loss/build.toml | 13 +- tdt-loss/tdt_loss/common.cuh | 87 +++++ tdt-loss/tdt_loss/lattice.cu | 234 +++++++++++ tdt-loss/tdt_loss/logprobs.cu | 391 +++++++++++++++++++ tdt-loss/tests/__init__.py | 0 tdt-loss/tests/test_tdt_loss.py | 490 +++++++++++++++--------- tdt-loss/torch-ext/tdt_loss/__init__.py | 229 ++++++----- tdt-loss/torch-ext/torch_binding.cpp | 24 +- tdt-loss/torch-ext/torch_binding.h | 52 +-- 9 files changed, 1190 insertions(+), 330 deletions(-) create mode 100644 tdt-loss/tdt_loss/common.cuh create mode 100644 tdt-loss/tdt_loss/lattice.cu create mode 100644 tdt-loss/tdt_loss/logprobs.cu create mode 100644 tdt-loss/tests/__init__.py diff --git a/tdt-loss/build.toml b/tdt-loss/build.toml index bb3444af3..f83977d85 100644 --- a/tdt-loss/build.toml +++ b/tdt-loss/build.toml @@ -15,16 +15,9 @@ src = [ [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", + "tdt_loss/common.cuh", + "tdt_loss/lattice.cu", + "tdt_loss/logprobs.cu", ] diff --git a/tdt-loss/tdt_loss/common.cuh b/tdt-loss/tdt_loss/common.cuh new file mode 100644 index 000000000..e593225d4 --- /dev/null +++ b/tdt-loss/tdt_loss/common.cuh @@ -0,0 +1,87 @@ +#pragma once + +#include + +#include +#include + +namespace tdt_loss { + +constexpr int kWarpSize = 32; +// Upper bound on the number of durations, only used to size shared memory. +constexpr int kMaxDurations = 64; + +__device__ __forceinline__ float neg_inf() { return -INFINITY; } + +// log(exp(a) + exp(b)), safe when either argument is -inf. +__device__ __forceinline__ float log_add(float a, float b) { + if (a == -INFINITY) return b; + if (b == -INFINITY) return a; + const float m = fmaxf(a, b); + return m + log1pf(expf(-fabsf(a - b))); +} + +// Merge two partial (max, sum(exp(x - max))) pairs of an online softmax. +__device__ __forceinline__ void merge_max_sum(float &m, float &s, float m_other, + float s_other) { + const float m_new = fmaxf(m, m_other); + if (m_new == -INFINITY) { + m = m_new; + s = 0.f; + return; + } + s = s * expf(m - m_new) + s_other * expf(m_other - m_new); + m = m_new; +} + +// Block-wide reduction of (max, sum) pairs. The result is valid in all threads. +// `shm_m` and `shm_s` must hold at least blockDim.x / kWarpSize floats. +__device__ __forceinline__ void block_reduce_max_sum(float &m, float &s, + float *shm_m, + float *shm_s) { + const int lane = threadIdx.x % kWarpSize; + const int warp = threadIdx.x / kWarpSize; + const int num_warps = (blockDim.x + kWarpSize - 1) / kWarpSize; + +#pragma unroll + for (int offset = kWarpSize / 2; offset > 0; offset /= 2) { + const float m_other = __shfl_xor_sync(0xffffffff, m, offset); + const float s_other = __shfl_xor_sync(0xffffffff, s, offset); + merge_max_sum(m, s, m_other, s_other); + } + if (lane == 0) { + shm_m[warp] = m; + shm_s[warp] = s; + } + __syncthreads(); + if (warp == 0) { + m = lane < num_warps ? shm_m[lane] : -INFINITY; + s = lane < num_warps ? shm_s[lane] : 0.f; +#pragma unroll + for (int offset = kWarpSize / 2; offset > 0; offset /= 2) { + const float m_other = __shfl_xor_sync(0xffffffff, m, offset); + const float s_other = __shfl_xor_sync(0xffffffff, s, offset); + merge_max_sum(m, s, m_other, s_other); + } + if (lane == 0) { + shm_m[0] = m; + shm_s[0] = s; + } + } + __syncthreads(); + m = shm_m[0]; + s = shm_s[0]; +} + +// Clamp the per-sample lengths to the padded lattice so that malformed inputs +// cannot cause out-of-bounds accesses. Returns T (frames) and U (labels); the +// lattice for the sample spans t in [0, T) and u in [0, U]. +__device__ __forceinline__ void sample_lengths(const int *source_lengths, + const int *target_lengths, + int b, int max_T, int max_U, + int &T, int &U) { + T = min(max(source_lengths[b], 0), max_T); + U = min(max(target_lengths[b], 0), max_U - 1); +} + +} // namespace tdt_loss diff --git a/tdt-loss/tdt_loss/lattice.cu b/tdt-loss/tdt_loss/lattice.cu new file mode 100644 index 000000000..23de82232 --- /dev/null +++ b/tdt-loss/tdt_loss/lattice.cu @@ -0,0 +1,234 @@ +// Forward (alpha) and backward (beta) recursions over the TDT lattice. +// +// For a sample with T frames and U labels the lattice nodes are (t, u) with +// t in [0, T) and u in [0, U]. From node (t, u), for every duration d: +// - a blank arc (only d > 0) goes to (t + d, u), +// - a label arc goes to (t + d, u + 1). +// The alignment ends with a blank arc from (T - d, U) that exits the lattice. +// +// Every arc strictly increases t + u, so all nodes of an anti-diagonal +// t + u = n are independent. One thread block handles one sample and sweeps +// the anti-diagonals, with the threads striding over u. + +#include +#include +#include + +#include "common.cuh" + +namespace tdt_loss { + +namespace { + +__device__ __forceinline__ void load_durations(const int *durations, int D, + int *shm_durations) { + for (int i = threadIdx.x; i < D; i += blockDim.x) { + shm_durations[i] = durations[i]; + } + __syncthreads(); +} + +__global__ void tdt_alpha_kernel(const float *__restrict__ blank_lp, + const float *__restrict__ label_lp, + const float *__restrict__ dur_lp, + const int *__restrict__ source_lengths, + const int *__restrict__ target_lengths, + const int *__restrict__ durations, int max_T, + int max_U, int D, float *__restrict__ alphas, + float *__restrict__ log_ll) { + __shared__ int shm_durations[kMaxDurations]; + + const int b = blockIdx.x; + int T, U; + sample_lengths(source_lengths, target_lengths, b, max_T, max_U, T, U); + load_durations(durations, D, shm_durations); + + const int64_t sample = static_cast(b) * max_T * max_U; + const float *blank = blank_lp + sample; + const float *label = label_lp + sample; + const float *dur = dur_lp + sample * D; + float *alpha = alphas + sample; + + for (int n = 0; n < T + U; ++n) { + const int u_min = max(0, n - T + 1); + const int u_max = min(n, U); + for (int u = u_min + threadIdx.x; u <= u_max; u += blockDim.x) { + const int t = n - u; + float a = n == 0 ? 0.f : neg_inf(); + for (int i = 0; i < D && n > 0; ++i) { + const int d = shm_durations[i]; + const int t_prev = t - d; + if (t_prev < 0) continue; + if (d > 0) { + const int64_t src = static_cast(t_prev) * max_U + u; + a = log_add(a, alpha[src] + blank[src] + dur[src * D + i]); + } + if (u > 0) { + const int64_t src = static_cast(t_prev) * max_U + u - 1; + a = log_add(a, alpha[src] + label[src] + dur[src * D + i]); + } + } + alpha[static_cast(t) * max_U + u] = a; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + float ll = neg_inf(); + for (int i = 0; i < D; ++i) { + const int d = shm_durations[i]; + const int t_prev = T - d; + if (d == 0 || t_prev < 0) continue; + const int64_t src = static_cast(t_prev) * max_U + U; + ll = log_add(ll, alpha[src] + blank[src] + dur[src * D + i]); + } + log_ll[b] = ll; + } +} + +__global__ void tdt_beta_kernel(const float *__restrict__ blank_lp, + const float *__restrict__ label_lp, + const float *__restrict__ dur_lp, + const int *__restrict__ source_lengths, + const int *__restrict__ target_lengths, + const int *__restrict__ durations, int max_T, + int max_U, int D, float *__restrict__ betas, + float *__restrict__ log_ll) { + __shared__ int shm_durations[kMaxDurations]; + + const int b = blockIdx.x; + int T, U; + sample_lengths(source_lengths, target_lengths, b, max_T, max_U, T, U); + load_durations(durations, D, shm_durations); + + const int64_t sample = static_cast(b) * max_T * max_U; + const float *blank = blank_lp + sample; + const float *label = label_lp + sample; + const float *dur = dur_lp + sample * D; + float *beta = betas + sample; + + for (int n = T + U - 1; n >= 0; --n) { + const int u_min = max(0, n - T + 1); + const int u_max = min(n, U); + for (int u = u_min + threadIdx.x; u <= u_max; u += blockDim.x) { + const int t = n - u; + const int64_t node = static_cast(t) * max_U + u; + float bt = neg_inf(); + for (int i = 0; i < D; ++i) { + const int d = shm_durations[i]; + const int t_next = t + d; + const float lp_dur = dur[node * D + i]; + if (d > 0) { + if (t_next < T) { + const int64_t dst = static_cast(t_next) * max_U + u; + bt = log_add(bt, beta[dst] + blank[node] + lp_dur); + } else if (t_next == T && u == U) { + bt = log_add(bt, blank[node] + lp_dur); + } + } + if (u < U && t_next < T) { + const int64_t dst = static_cast(t_next) * max_U + u + 1; + bt = log_add(bt, beta[dst] + label[node] + lp_dur); + } + } + beta[node] = bt; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + log_ll[b] = T > 0 ? beta[0] : neg_inf(); + } +} + +int lattice_block_size(int64_t max_U) { + int threads = 32; + while (threads < 1024 && threads < max_U) threads *= 2; + return threads; +} + +void check_lattice_args(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 const &out, torch::Tensor const &log_ll) { + TORCH_CHECK(blank_lp.is_cuda(), "blank_lp must be a CUDA tensor"); + TORCH_CHECK(blank_lp.dim() == 3, "blank_lp must have shape (batch, T, U+1)"); + for (auto const *x : {&blank_lp, &label_lp, &dur_lp, &out, &log_ll}) { + TORCH_CHECK(x->device() == blank_lp.device(), + "all tensors must be on the same device"); + TORCH_CHECK(x->scalar_type() == torch::kFloat32, + "lattice tensors must be float32"); + TORCH_CHECK(x->is_contiguous(), "lattice tensors must be contiguous"); + } + for (auto const *x : {&source_lengths, &target_lengths, &durations}) { + TORCH_CHECK(x->device() == blank_lp.device(), + "all tensors must be on the same device"); + TORCH_CHECK(x->scalar_type() == torch::kInt32, + "lengths and durations must be int32"); + TORCH_CHECK(x->is_contiguous(), "lengths and durations must be contiguous"); + } + const int64_t B = blank_lp.size(0); + const int64_t nodes = blank_lp.numel(); + TORCH_CHECK(label_lp.numel() == nodes && out.numel() == nodes, + "lattice tensors must have shape (batch, T, U+1)"); + TORCH_CHECK(dur_lp.numel() == nodes * durations.numel(), + "dur_lp must have shape (batch, T, U+1, num_durations)"); + TORCH_CHECK(source_lengths.numel() == B && target_lengths.numel() == B && + log_ll.numel() == B, + "lengths and log_ll must have shape (batch,)"); + TORCH_CHECK(durations.numel() <= kMaxDurations, "at most ", kMaxDurations, + " durations are supported"); +} + +} // namespace + +} // namespace tdt_loss + +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) { + using namespace tdt_loss; + check_lattice_args(blank_lp, label_lp, dur_lp, source_lengths, + target_lengths, durations, alphas, log_ll); + const int64_t B = blank_lp.size(0); + if (B == 0) return; + + const at::cuda::OptionalCUDAGuard device_guard(device_of(blank_lp)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + tdt_alpha_kernel<<>>( + blank_lp.data_ptr(), label_lp.data_ptr(), + dur_lp.data_ptr(), source_lengths.data_ptr(), + target_lengths.data_ptr(), durations.data_ptr(), + blank_lp.size(1), blank_lp.size(2), durations.numel(), + alphas.data_ptr(), log_ll.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +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 &log_ll) { + using namespace tdt_loss; + check_lattice_args(blank_lp, label_lp, dur_lp, source_lengths, + target_lengths, durations, betas, log_ll); + const int64_t B = blank_lp.size(0); + if (B == 0) return; + + const at::cuda::OptionalCUDAGuard device_guard(device_of(blank_lp)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + tdt_beta_kernel<<>>( + blank_lp.data_ptr(), label_lp.data_ptr(), + dur_lp.data_ptr(), source_lengths.data_ptr(), + target_lengths.data_ptr(), durations.data_ptr(), + blank_lp.size(1), blank_lp.size(2), durations.numel(), + betas.data_ptr(), log_ll.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/tdt-loss/tdt_loss/logprobs.cu b/tdt-loss/tdt_loss/logprobs.cu new file mode 100644 index 000000000..2f93202c6 --- /dev/null +++ b/tdt-loss/tdt_loss/logprobs.cu @@ -0,0 +1,391 @@ +// Per-lattice-node kernels of the TDT loss: the log-softmax gather that feeds +// the alpha/beta recursions, and the fused gradient w.r.t. the logits. +// +// Both kernels use one thread block per lattice node (b, t, u) and stream the +// vocabulary dimension. Nodes outside of a sample's lattice (t >= T or u > U) +// never read the logits. + +#include +#include +#include + +#include + +#include "common.cuh" + +namespace tdt_loss { + +namespace { + +template +__global__ void tdt_logprobs_fwd_kernel( + const scalar_t *__restrict__ token_logits, int64_t tok_sb, int64_t tok_st, + int64_t tok_su, const scalar_t *__restrict__ duration_logits, + int64_t dur_sb, int64_t dur_st, int64_t dur_su, + const int *__restrict__ targets, int64_t targets_stride, + const int *__restrict__ source_lengths, + const int *__restrict__ target_lengths, int max_T, int max_U, int V, int D, + int blank_id, float sigma, float *__restrict__ blank_lp, + float *__restrict__ label_lp, float *__restrict__ dur_lp, + float *__restrict__ token_lse) { + __shared__ float shm_m[kWarpSize]; + __shared__ float shm_s[kWarpSize]; + + const int64_t row = blockIdx.x; + const int u = row % max_U; + const int t = (row / max_U) % max_T; + const int b = row / (static_cast(max_U) * max_T); + + int T, U; + sample_lengths(source_lengths, target_lengths, b, max_T, max_U, T, U); + + if (t >= T || u > U) { + if (threadIdx.x == 0) { + blank_lp[row] = neg_inf(); + label_lp[row] = neg_inf(); + token_lse[row] = 0.f; + } + for (int i = threadIdx.x; i < D; i += blockDim.x) { + dur_lp[row * D + i] = neg_inf(); + } + return; + } + + const scalar_t *x = token_logits + b * tok_sb + t * tok_st + u * tok_su; + + // Online softmax normalizer over the vocabulary. + float m = neg_inf(); + float s = 0.f; + for (int v = threadIdx.x; v < V; v += blockDim.x) { + const float xv = static_cast(x[v]); + if (xv > m) { + s = s * expf(m - xv) + 1.f; + m = xv; + } else if (xv != -INFINITY) { + s += expf(xv - m); + } + } + block_reduce_max_sum(m, s, shm_m, shm_s); + + if (threadIdx.x == 0) { + const float lse = m + logf(s); + token_lse[row] = lse; + blank_lp[row] = static_cast(x[blank_id]) - lse - sigma; + + float label = neg_inf(); + if (u < U) { + const int target = targets[b * targets_stride + u]; + if (target >= 0 && target < V) { + label = static_cast(x[target]) - lse - sigma; + } + } + label_lp[row] = label; + + // The duration head is tiny (a handful of classes), a serial + // log-softmax is cheaper than another block reduction. + const scalar_t *y = + duration_logits + b * dur_sb + t * dur_st + u * dur_su; + float dm = neg_inf(); + for (int i = 0; i < D; ++i) dm = fmaxf(dm, static_cast(y[i])); + float ds = 0.f; + for (int i = 0; i < D; ++i) ds += expf(static_cast(y[i]) - dm); + const float dlse = dm + logf(ds); + for (int i = 0; i < D; ++i) { + dur_lp[row * D + i] = static_cast(y[i]) - dlse; + } + } +} + +template +__global__ void tdt_logits_grad_kernel( + const scalar_t *__restrict__ token_logits, int64_t tok_sb, int64_t tok_st, + int64_t tok_su, const int *__restrict__ targets, int64_t targets_stride, + const int *__restrict__ source_lengths, + const int *__restrict__ target_lengths, const int *__restrict__ durations, + const float *__restrict__ blank_lp, const float *__restrict__ label_lp, + const float *__restrict__ dur_lp, const float *__restrict__ token_lse, + const float *__restrict__ alphas, const float *__restrict__ betas, + const float *__restrict__ log_ll, const float *__restrict__ grad_loss, + int max_T, int max_U, int V, int D, int blank_id, + scalar_t *__restrict__ grad_token_logits, + scalar_t *__restrict__ grad_duration_logits) { + // Posterior occupancy of the blank / label arc leaving this node, for each + // duration. + __shared__ float shm_blank[kMaxDurations]; + __shared__ float shm_label[kMaxDurations]; + + const int64_t row = blockIdx.x; + const int u = row % max_U; + const int t = (row / max_U) % max_T; + const int b = row / (static_cast(max_U) * max_T); + + int T, U; + sample_lengths(source_lengths, target_lengths, b, max_T, max_U, T, U); + + scalar_t *gx = grad_token_logits + row * V; + scalar_t *gd = grad_duration_logits + row * D; + const scalar_t zero = static_cast(0.f); + + const float ll = log_ll[b]; + // Infeasible samples (infinite loss) get a zero gradient instead of NaNs. + if (t >= T || u > U || !isfinite(ll)) { + for (int v = threadIdx.x; v < V; v += blockDim.x) gx[v] = zero; + for (int i = threadIdx.x; i < D; i += blockDim.x) gd[i] = zero; + return; + } + + const int64_t sample = static_cast(b) * max_T * max_U; + const float alpha = alphas[row]; + const float blank = blank_lp[row]; + const float label = label_lp[row]; + + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const int d = durations[i]; + const int t_next = t + d; + const float lp_dur = dur_lp[row * D + i]; + + float p_blank = 0.f; + if (d > 0) { + float beta_next = neg_inf(); + if (t_next < T) { + beta_next = betas[sample + static_cast(t_next) * max_U + u]; + } else if (t_next == T && u == U) { + beta_next = 0.f; // terminal arc + } + if (beta_next != -INFINITY) { + p_blank = expf(alpha + blank + lp_dur + beta_next - ll); + } + } + + float p_label = 0.f; + if (u < U && t_next < T) { + const float beta_next = + betas[sample + static_cast(t_next) * max_U + u + 1]; + p_label = expf(alpha + label + lp_dur + beta_next - ll); + } + + shm_blank[i] = p_blank; + shm_label[i] = p_label; + } + __syncthreads(); + + float occ_blank = 0.f; + float occ_label = 0.f; + for (int i = 0; i < D; ++i) { + occ_blank += shm_blank[i]; + occ_label += shm_label[i]; + } + + // loss = -log_ll, so d loss / d log_prob(arc) = -occupancy(arc). + const float go = grad_loss[b]; + const float g_blank = -go * occ_blank; + const float g_label = -go * occ_label; + const float g_total = g_blank + g_label; + + // Through log_softmax: d/dx_j = g_j - softmax_j * sum_k g_k. + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float g_dur = -go * (shm_blank[i] + shm_label[i]); + gd[i] = static_cast(g_dur - expf(dur_lp[row * D + i]) * g_total); + } + + if (g_blank == 0.f && g_label == 0.f) { + for (int v = threadIdx.x; v < V; v += blockDim.x) gx[v] = zero; + return; + } + + const scalar_t *x = token_logits + b * tok_sb + t * tok_st + u * tok_su; + const float lse = token_lse[row]; + const int target = u < U ? targets[b * targets_stride + u] : -1; + for (int v = threadIdx.x; v < V; v += blockDim.x) { + float g = -expf(static_cast(x[v]) - lse) * g_total; + if (v == blank_id) g += g_blank; + if (v == target) g += g_label; + gx[v] = static_cast(g); + } +} + +int row_block_size(int64_t V) { + int threads = 32; + while (threads < 256 && threads < V) threads *= 2; + return threads; +} + +void check_logits(torch::Tensor const &token_logits, + torch::Tensor const &duration_logits) { + TORCH_CHECK(token_logits.is_cuda(), "token_logits must be a CUDA tensor"); + TORCH_CHECK(token_logits.dim() == 4, + "token_logits must have shape (batch, T, U+1, vocab_size+1)"); + TORCH_CHECK(duration_logits.dim() == 4, + "duration_logits must have shape (batch, T, U+1, num_durations)"); + TORCH_CHECK(duration_logits.device() == token_logits.device(), + "token_logits and duration_logits must be on the same device"); + TORCH_CHECK(duration_logits.scalar_type() == token_logits.scalar_type(), + "token_logits and duration_logits must have the same dtype"); + for (int i = 0; i < 3; ++i) { + TORCH_CHECK(duration_logits.size(i) == token_logits.size(i), + "token_logits and duration_logits must agree on (batch, T, " + "U+1)"); + } + TORCH_CHECK(token_logits.stride(3) == 1 && duration_logits.stride(3) == 1, + "the last dimension of the logits must be contiguous"); + TORCH_CHECK(duration_logits.size(3) <= kMaxDurations, + "at most ", kMaxDurations, " durations are supported"); +} + +void check_int_tensor(torch::Tensor const &x, char const *name, + torch::Tensor const &like) { + TORCH_CHECK(x.device() == like.device(), name, + " must be on the same device as the logits"); + TORCH_CHECK(x.scalar_type() == torch::kInt32, name, " must be int32"); + TORCH_CHECK(x.is_contiguous(), name, " must be contiguous"); +} + +void check_float_out(torch::Tensor const &x, char const *name, + torch::Tensor const &like) { + TORCH_CHECK(x.device() == like.device(), name, + " must be on the same device as the logits"); + TORCH_CHECK(x.scalar_type() == torch::kFloat32, name, " must be float32"); + TORCH_CHECK(x.is_contiguous(), name, " must be contiguous"); +} + +} // namespace + +} // namespace tdt_loss + +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, + torch::Tensor &token_lse) { + using namespace tdt_loss; + + check_logits(token_logits, duration_logits); + const int64_t B = token_logits.size(0); + const int64_t max_T = token_logits.size(1); + const int64_t max_U = token_logits.size(2); + const int64_t V = token_logits.size(3); + const int64_t D = duration_logits.size(3); + + TORCH_CHECK(blank_id >= 0 && blank_id < V, "blank_id out of range"); + TORCH_CHECK(targets.dim() == 2 && targets.size(0) == B && + targets.size(1) >= max_U - 1, + "targets must have shape (batch, U)"); + check_int_tensor(targets, "targets", token_logits); + check_int_tensor(source_lengths, "source_lengths", token_logits); + check_int_tensor(target_lengths, "target_lengths", token_logits); + TORCH_CHECK(source_lengths.numel() == B && target_lengths.numel() == B, + "lengths must have shape (batch,)"); + for (auto *out : {&blank_lp, &label_lp, &token_lse}) { + check_float_out(*out, "log-prob output", token_logits); + TORCH_CHECK(out->numel() == B * max_T * max_U, + "log-prob outputs must have shape (batch, T, U+1)"); + } + check_float_out(dur_lp, "dur_lp", token_logits); + TORCH_CHECK(dur_lp.numel() == B * max_T * max_U * D, + "dur_lp must have shape (batch, T, U+1, num_durations)"); + + const int64_t rows = B * max_T * max_U; + if (rows == 0) return; + TORCH_CHECK(rows <= std::numeric_limits::max(), + "batch * T * (U+1) is too large"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(token_logits)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const int threads = row_block_size(V); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, token_logits.scalar_type(), "tdt_logprobs_fwd", + [&] { + tdt_logprobs_fwd_kernel<<(rows), threads, 0, stream>>>( + token_logits.data_ptr(), token_logits.stride(0), + token_logits.stride(1), token_logits.stride(2), + duration_logits.data_ptr(), duration_logits.stride(0), + duration_logits.stride(1), duration_logits.stride(2), + targets.data_ptr(), targets.stride(0), + source_lengths.data_ptr(), target_lengths.data_ptr(), + max_T, max_U, V, D, blank_id, static_cast(sigma), + blank_lp.data_ptr(), label_lp.data_ptr(), + dur_lp.data_ptr(), token_lse.data_ptr()); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void tdt_logits_grad( + torch::Tensor const &token_logits, torch::Tensor const &targets, + torch::Tensor const &source_lengths, torch::Tensor const &target_lengths, + torch::Tensor const &durations, torch::Tensor const &blank_lp, + torch::Tensor const &label_lp, torch::Tensor const &dur_lp, + torch::Tensor const &token_lse, torch::Tensor const &alphas, + torch::Tensor const &betas, torch::Tensor const &log_ll, + torch::Tensor const &grad_loss, int64_t blank_id, + torch::Tensor &grad_token_logits, torch::Tensor &grad_duration_logits) { + using namespace tdt_loss; + + TORCH_CHECK(token_logits.is_cuda() && token_logits.dim() == 4, + "token_logits must be a 4D CUDA tensor"); + TORCH_CHECK(token_logits.stride(3) == 1, + "the last dimension of token_logits must be contiguous"); + const int64_t B = token_logits.size(0); + const int64_t max_T = token_logits.size(1); + const int64_t max_U = token_logits.size(2); + const int64_t V = token_logits.size(3); + const int64_t D = durations.numel(); + + TORCH_CHECK(D <= kMaxDurations, "at most ", kMaxDurations, + " durations are supported"); + TORCH_CHECK(blank_id >= 0 && blank_id < V, "blank_id out of range"); + check_int_tensor(targets, "targets", token_logits); + check_int_tensor(source_lengths, "source_lengths", token_logits); + check_int_tensor(target_lengths, "target_lengths", token_logits); + check_int_tensor(durations, "durations", token_logits); + for (auto *in : {&blank_lp, &label_lp, &token_lse, &alphas, &betas}) { + check_float_out(*in, "lattice tensor", token_logits); + TORCH_CHECK(in->numel() == B * max_T * max_U, + "lattice tensors must have shape (batch, T, U+1)"); + } + check_float_out(dur_lp, "dur_lp", token_logits); + check_float_out(log_ll, "log_ll", token_logits); + check_float_out(grad_loss, "grad_loss", token_logits); + TORCH_CHECK(log_ll.numel() == B && grad_loss.numel() == B, + "log_ll and grad_loss must have shape (batch,)"); + TORCH_CHECK(grad_token_logits.is_contiguous() && + grad_token_logits.sizes() == token_logits.sizes() && + grad_token_logits.scalar_type() == token_logits.scalar_type(), + "grad_token_logits must be contiguous and match token_logits"); + TORCH_CHECK(grad_duration_logits.is_contiguous() && + grad_duration_logits.numel() == B * max_T * max_U * D && + grad_duration_logits.scalar_type() == + token_logits.scalar_type(), + "grad_duration_logits must be contiguous and match " + "duration_logits"); + + const int64_t rows = B * max_T * max_U; + if (rows == 0) return; + TORCH_CHECK(rows <= std::numeric_limits::max(), + "batch * T * (U+1) is too large"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(token_logits)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const int threads = row_block_size(V); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, token_logits.scalar_type(), "tdt_logits_grad", + [&] { + tdt_logits_grad_kernel<<(rows), threads, 0, stream>>>( + token_logits.data_ptr(), token_logits.stride(0), + token_logits.stride(1), token_logits.stride(2), + targets.data_ptr(), targets.stride(0), + source_lengths.data_ptr(), target_lengths.data_ptr(), + durations.data_ptr(), blank_lp.data_ptr(), + label_lp.data_ptr(), dur_lp.data_ptr(), + token_lse.data_ptr(), alphas.data_ptr(), + betas.data_ptr(), log_ll.data_ptr(), + grad_loss.data_ptr(), max_T, max_U, V, D, blank_id, + grad_token_logits.data_ptr(), + grad_duration_logits.data_ptr()); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/tdt-loss/tests/__init__.py b/tdt-loss/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tdt-loss/tests/test_tdt_loss.py b/tdt-loss/tests/test_tdt_loss.py index 236982e6a..580fbffa6 100644 --- a/tdt-loss/tests/test_tdt_loss.py +++ b/tdt-loss/tests/test_tdt_loss.py @@ -1,4 +1,12 @@ -"""Tests for TDT loss kernel, validated against a naive loop-based reference.""" +"""Tests for the TDT loss kernel. + +The kernel is checked against two pure PyTorch references: +- `tdt_loss_reference`: the vectorized implementation used in transformers, whose autograd + gradients are used to validate the kernel backward; +- `tdt_loss_naive`: a loop-based implementation that is easy to verify by hand. +""" + +from pathlib import Path import pytest import torch @@ -7,215 +15,337 @@ 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 + import tdt_loss as tdt_loss_kernel +except ImportError: + from kernels import get_local_kernel - token_logits = token_logits.float() - duration_logits = duration_logits.float() + tdt_loss_kernel = get_local_kernel(repo_path=Path(__file__).parent.parent, package_name="tdt_loss") - tok_lp = torch.log_softmax(token_logits, dim=-1) - sigma - dur_lp = torch.log_softmax(duration_logits, dim=-1) +DEVICE = torch.device("cuda") +DURATIONS = [0, 1, 2, 3, 4] +REDUCTIONS = ["mean_volume", "mean_batch", "mean", "sum", "none"] - 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 +def reduce(losses, target_lengths, reduction): + target_lengths = target_lengths.float() + if reduction == "mean_volume": + return losses.sum() / target_lengths.sum() + if reduction == "mean_batch": + return losses.mean() + if reduction == "mean": + return (losses / target_lengths).mean() + if reduction == "sum": + return losses.sum() + return losses - 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) +def tdt_loss_reference( + token_logits, duration_logits, targets, logit_lengths, target_lengths, blank_id, durations, sigma=0.0 +): + """Per-sample TDT loss, as implemented in `transformers.loss.loss_tdt.tdt_loss`.""" + device = token_logits.device + batch_size, max_t, max_u, _ = token_logits.shape + + token_log_probs = torch.log_softmax(token_logits.float(), dim=-1) - sigma + duration_log_probs = torch.log_softmax(duration_logits.float(), dim=-1) + + log_alpha = torch.full((batch_size, max_t, max_u), float("-inf"), device=device) + log_alpha[:, 0, 0] = 0.0 + blank_log_probs = token_log_probs[:, :, :, blank_id] + if max_u > 1: + targets_expanded = targets.long().unsqueeze(1).expand(-1, max_t, -1) + label_log_probs = torch.gather( + token_log_probs[:, :, : max_u - 1, :], dim=3, index=targets_expanded.unsqueeze(-1) + ).squeeze(-1) + + neg_inf = torch.tensor(float("-inf"), device=device) + for n in range(1, max_t + max_u - 1): + u_indices = torch.arange(max(0, n - max_t + 1), min(n + 1, max_u), device=device) + t_indices = n - u_indices + candidates = [] for i, dur in enumerate(durations): - if dur == 0: + t_prev = t_indices - dur + valid_t = t_prev >= 0 + if not valid_t.any(): 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)) + t_src = t_prev.clamp(min=0) + if dur > 0: + contrib = ( + log_alpha[:, t_src, u_indices] + + blank_log_probs[:, t_src, u_indices] + + duration_log_probs[:, t_src, u_indices, i] + ) + candidates.append(torch.where(valid_t.unsqueeze(0), contrib, neg_inf)) + valid_both = valid_t & (u_indices > 0) + if valid_both.any(): + u_src = (u_indices - 1).clamp(min=0) + u_src_label = u_src.clamp(max=max_u - 2) if max_u > 1 else u_src + contrib = ( + log_alpha[:, t_src, u_src] + + label_log_probs[:, t_src, u_src_label] + + duration_log_probs[:, t_src, u_src, i] + ) + candidates.append(torch.where(valid_both.unsqueeze(0), contrib, neg_inf)) + if candidates: + log_alpha[:, t_indices, u_indices] = torch.logsumexp(torch.stack(candidates), dim=0) + + batch_idx = torch.arange(batch_size, device=device) + target_lengths = target_lengths.long() + log_probs = torch.full((batch_size,), float("-inf"), device=device) + for i, dur in enumerate(durations): + if dur == 0: + continue + t_final = logit_lengths.long() - dur + valid = t_final >= 0 + if not valid.any(): + continue + t_clamped = t_final.clamp(min=0) + terminal = ( + log_alpha[batch_idx, t_clamped, target_lengths] + + token_log_probs[batch_idx, t_clamped, target_lengths, blank_id] + + duration_log_probs[batch_idx, t_clamped, target_lengths, i] + ) + log_probs = torch.where(valid, torch.logsumexp(torch.stack([log_probs, terminal]), dim=0), log_probs) + return -log_probs - losses.append(-ll) - return torch.stack(losses) +def tdt_loss_naive(token_logits, duration_logits, targets, logit_lengths, target_lengths, blank_id, durations, sigma=0.0): + """Per-sample TDT loss with explicit loops over the lattice, in float64 on the CPU.""" + token_lp = (torch.log_softmax(token_logits.double().cpu(), dim=-1) - sigma).tolist() + dur_lp = torch.log_softmax(duration_logits.double().cpu(), dim=-1).tolist() + targets = targets.tolist() + neg_inf = float("-inf") + def log_add(a, b): + return torch.logaddexp(torch.tensor(a, dtype=torch.float64), torch.tensor(b, dtype=torch.float64)).item() + + losses = [] + for b in range(len(targets)): + T, U = int(logit_lengths[b]), int(target_lengths[b]) + alpha = [[neg_inf] * (U + 1) for _ in range(T)] + alpha[0][0] = 0.0 + for t in range(T): + for u in range(U + 1): + if t == 0 and u == 0: + continue + for i, d in enumerate(durations): + t_src = t - d + if t_src < 0: + continue + if d > 0: + arc = alpha[t_src][u] + token_lp[b][t_src][u][blank_id] + dur_lp[b][t_src][u][i] + alpha[t][u] = log_add(alpha[t][u], arc) + if u > 0: + label = targets[b][u - 1] + arc = alpha[t_src][u - 1] + token_lp[b][t_src][u - 1][label] + dur_lp[b][t_src][u - 1][i] + alpha[t][u] = log_add(alpha[t][u], arc) + log_ll = neg_inf + for i, d in enumerate(durations): + if d > 0 and T - d >= 0: + t_src = T - d + log_ll = log_add(log_ll, alpha[t_src][U] + token_lp[b][t_src][U][blank_id] + dur_lp[b][t_src][U][i]) + losses.append(-log_ll) + return torch.tensor(losses, dtype=torch.float64) + + +def make_inputs(B, T, U, V, durations, dtype=torch.float32, seed=0, joint=False): + """Random inputs with variable lengths. `joint=True` returns slices of a single joint tensor, + like `ParakeetForTDT` does.""" + generator = torch.Generator(device="cpu").manual_seed(seed) + D = len(durations) + logits = torch.randn(B, T, U + 1, V + D, generator=generator).to(DEVICE, dtype) + if joint: + token_logits, duration_logits = logits[..., :V], logits[..., V:] + else: + token_logits, duration_logits = logits[..., :V].contiguous(), logits[..., V:].contiguous() + targets = torch.randint(1, V, (B, U), generator=generator).to(DEVICE, torch.int32) + logit_lengths = torch.tensor([T - (b % 3) for b in range(B)], device=DEVICE, dtype=torch.int32) + target_lengths = torch.tensor([U - (b % 2) for b in range(B)], device=DEVICE, dtype=torch.int32) + return token_logits, duration_logits, targets, logit_lengths, target_lengths + + +def run_kernel(token_logits, duration_logits, targets, logit_lengths, target_lengths, durations, sigma, reduction): + return tdt_loss_kernel.tdt_loss( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + durations, + 0, + sigma=sigma, + reduction=reduction, + ) -# --------------------------------------------------------------------------- -# 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 +def test_matches_naive(sigma): + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs(3, 12, 5, 16, DURATIONS) + expected = tdt_loss_naive( + token_logits, duration_logits, targets, logit_lengths, target_lengths, 0, DURATIONS, sigma=sigma + ) + losses = run_kernel( + token_logits, duration_logits, targets, logit_lengths, target_lengths, DURATIONS, sigma, "none" ) + torch.testing.assert_close(losses.double().cpu(), expected, atol=1e-4, rtol=1e-5) - # 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, +@pytest.mark.parametrize("reduction", REDUCTIONS) +@pytest.mark.parametrize("sigma", [0.0, 0.05]) +@pytest.mark.parametrize("durations", [DURATIONS, [1, 2], [0, 1, 3, 5]]) +def test_forward_backward_match_reference(durations, sigma, reduction): + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs(4, 30, 8, 64, durations) + + ref_tok = token_logits.clone().requires_grad_(True) + ref_dur = duration_logits.clone().requires_grad_(True) + ref_losses = tdt_loss_reference( + ref_tok, ref_dur, targets, logit_lengths, target_lengths, 0, durations, sigma=sigma ) + expected = reduce(ref_losses, target_lengths, reduction) + expected.sum().backward() - if reduction == "none": - ref = ref_per_sample - elif reduction == "mean": - ref = (ref_per_sample / target_lengths.float()).mean() - else: - ref = ref_per_sample.sum() + tok = token_logits.clone().requires_grad_(True) + dur = duration_logits.clone().requires_grad_(True) + loss = run_kernel(tok, dur, targets, logit_lengths, target_lengths, durations, sigma, reduction) + loss.sum().backward() - torch.testing.assert_close(kernel_out.cpu(), ref.cpu(), atol=1e-4, rtol=1e-4) + torch.testing.assert_close(loss, expected, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(tok.grad, ref_tok.grad, atol=1e-5, rtol=1e-4) + torch.testing.assert_close(dur.grad, ref_dur.grad, atol=1e-5, 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") +def test_joint_logits_slices(): + """Non-contiguous slices of a joint output (as in ParakeetForTDT) are handled without a copy.""" + V = 128 + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs( + 3, 25, 6, V, DURATIONS, joint=True + ) + assert not token_logits.is_contiguous() and not duration_logits.is_contiguous() - 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", + joint = token_logits._base.clone().requires_grad_(True) + loss = run_kernel( + joint[..., :V], joint[..., V:], targets, logit_lengths, target_lengths, DURATIONS, 0.0, "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", + loss.backward() + + ref_joint = token_logits._base.clone().requires_grad_(True) + expected = reduce( + tdt_loss_reference( + ref_joint[..., :V], ref_joint[..., V:], targets, logit_lengths, target_lengths, 0, DURATIONS + ), + target_lengths, + "mean", ) - loss2.backward() - - # Losses should match - torch.testing.assert_close(loss1, loss2) + expected.backward() - # 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) + torch.testing.assert_close(loss, expected, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(joint.grad, ref_joint.grad, atol=1e-5, rtol=1e-4) -def test_tdt_loss_batch_size_one(): - """Sanity check with batch size 1.""" - torch.manual_seed(0) - device = torch.device("cuda") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_half_precision(dtype): + """Half-precision logits are read directly and gradients come back in the input dtype.""" + if dtype == torch.bfloat16 and torch.cuda.get_device_capability() < (8, 0): + pytest.skip("bfloat16 requires compute capability >= 8.0") + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs( + 4, 30, 8, 256, DURATIONS, dtype=dtype + ) - 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", + tok = token_logits.clone().requires_grad_(True) + dur = duration_logits.clone().requires_grad_(True) + loss = run_kernel(tok, dur, targets, logit_lengths, target_lengths, DURATIONS, 0.0, "mean") + loss.backward() + assert tok.grad.dtype == dtype and dur.grad.dtype == dtype + + # The reference upcasts the same half-precision values to float32. + ref_tok = token_logits.float().requires_grad_(True) + ref_dur = duration_logits.float().requires_grad_(True) + expected = reduce( + tdt_loss_reference(ref_tok, ref_dur, targets, logit_lengths, target_lengths, 0, DURATIONS), + target_lengths, + "mean", ) - torch.testing.assert_close(out.cpu(), ref.cpu(), atol=1e-4, rtol=1e-4) + expected.backward() + + torch.testing.assert_close(loss, expected, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(tok.grad.float(), ref_tok.grad, atol=1e-3, rtol=1e-2) + torch.testing.assert_close(dur.grad.float(), ref_dur.grad, atol=1e-3, rtol=1e-2) -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 +def test_parakeet_like_shapes(): + """Vocabulary and label lengths of the size used by Parakeet TDT.""" + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs( + 2, 60, 20, 8193, DURATIONS, joint=True + ) + tok = token_logits.clone().requires_grad_(True) + dur = duration_logits.clone().requires_grad_(True) + losses = run_kernel(tok, dur, targets, logit_lengths, target_lengths, DURATIONS, 0.0, "none") + losses.sum().backward() + + ref_tok = token_logits.clone().requires_grad_(True) + ref_dur = duration_logits.clone().requires_grad_(True) + expected = tdt_loss_reference(ref_tok, ref_dur, targets, logit_lengths, target_lengths, 0, DURATIONS) + expected.sum().backward() + + torch.testing.assert_close(losses, expected, atol=1e-3, rtol=1e-4) + torch.testing.assert_close(tok.grad, ref_tok.grad, atol=1e-5, rtol=1e-4) + torch.testing.assert_close(dur.grad, ref_dur.grad, atol=1e-5, rtol=1e-4) + + +def test_long_targets(): + """More labels than threads in a block (the lattice kernels stride over u).""" + U = 1100 + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs(1, 6, U, 8, DURATIONS) + tok = token_logits.clone().requires_grad_(True) + loss = run_kernel(tok, duration_logits, targets, logit_lengths, target_lengths, DURATIONS, 0.0, "sum") + loss.backward() + + ref_tok = token_logits.clone().requires_grad_(True) + expected = tdt_loss_reference(ref_tok, duration_logits, targets, logit_lengths, target_lengths, 0, DURATIONS) + expected.sum().backward() + + torch.testing.assert_close(loss, expected.sum(), atol=1e-2, rtol=1e-4) + torch.testing.assert_close(tok.grad, ref_tok.grad, atol=1e-5, rtol=1e-4) + + +def test_empty_targets(): + """Samples without labels: the only path is a sequence of blanks.""" + token_logits, duration_logits, _, logit_lengths, _ = make_inputs(2, 10, 0, 16, DURATIONS) + targets = torch.empty(2, 0, device=DEVICE, dtype=torch.int32) + target_lengths = torch.zeros(2, device=DEVICE, dtype=torch.int32) + + losses = run_kernel(token_logits, duration_logits, targets, logit_lengths, target_lengths, DURATIONS, 0.0, "none") + expected = tdt_loss_naive(token_logits, duration_logits, targets, logit_lengths, target_lengths, 0, DURATIONS) + torch.testing.assert_close(losses.double().cpu(), expected, atol=1e-4, rtol=1e-5) + + +def test_infeasible_alignment(): + """No valid path: the loss is infinite and the gradient is zero rather than NaN.""" + durations = [0, 4] + token_logits, duration_logits, targets, _, target_lengths = make_inputs(2, 6, 3, 16, durations) + # Sample 0 emits all labels at t=0 and exits with a blank of duration 4; sample 1 has T=3 < 4. + logit_lengths = torch.tensor([4, 3], device=DEVICE, dtype=torch.int32) + + tok = token_logits.clone().requires_grad_(True) + losses = run_kernel(tok, duration_logits, targets, logit_lengths, target_lengths, durations, 0.0, "none") + assert torch.isfinite(losses[0]) and torch.isinf(losses[1]) + losses.backward(torch.ones_like(losses)) + assert torch.isfinite(tok.grad).all() + assert tok.grad[0].abs().sum() > 0 + assert (tok.grad[1] == 0).all() + + +def test_deterministic(): + inputs = make_inputs(4, 30, 8, 64, DURATIONS) + grads = [] + for _ in range(2): + tok = inputs[0].clone().requires_grad_(True) + run_kernel(tok, *inputs[1:], DURATIONS, 0.0, "mean").backward() + grads.append(tok.grad) + assert torch.equal(grads[0], grads[1]) + + +def test_invalid_reduction(): + inputs = make_inputs(1, 5, 2, 4, [0, 1]) 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", - ) + run_kernel(*inputs, [0, 1], 0.0, "invalid") diff --git a/tdt-loss/torch-ext/tdt_loss/__init__.py b/tdt-loss/torch-ext/tdt_loss/__init__.py index 8649c7bbe..07b026a97 100644 --- a/tdt-loss/torch-ext/tdt_loss/__init__.py +++ b/tdt-loss/torch-ext/tdt_loss/__init__.py @@ -1,14 +1,24 @@ """TDT (Token-and-Duration Transducer) loss CUDA kernel.""" -from typing import List, Union +from typing import Sequence, Union import torch from ._ops import ops +__all__ = ["tdt_loss"] + +_REDUCTIONS = ("mean_volume", "mean_batch", "mean", "sum", "none") + + +def _last_dim_contiguous(x: torch.Tensor) -> torch.Tensor: + # The kernels accept arbitrary strides for the (batch, T, U) dims, so slices of a + # joint `(..., vocab_size + num_durations)` output can be passed without a copy. + return x if x.stride(-1) == 1 else x.contiguous() + class TDTLoss(torch.autograd.Function): - """Custom autograd function for TDT loss.""" + """Per-sample TDT loss (negative log-likelihood) with a CUDA forward and backward.""" @staticmethod def forward( @@ -20,85 +30,97 @@ def forward( target_lengths: torch.Tensor, durations: torch.Tensor, blank_id: int, - sigma: float = 0.0, + sigma: float, ) -> 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) + token_logits = _last_dim_contiguous(token_logits) + duration_logits = _last_dim_contiguous(duration_logits) + B, max_T, max_U, _ = token_logits.shape + D = duration_logits.shape[-1] + f32 = dict(device=token_logits.device, dtype=torch.float32) + + blank_lp = torch.empty(B, max_T, max_U, **f32) + label_lp = torch.empty(B, max_T, max_U, **f32) + dur_lp = torch.empty(B, max_T, max_U, D, **f32) + token_lse = torch.empty(B, max_T, max_U, **f32) 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, + token_logits, + duration_logits, + targets, + source_lengths, + target_lengths, + blank_id, + sigma, + blank_lp, + label_lp, + dur_lp, + token_lse, ) + alphas = torch.empty(B, max_T, max_U, **f32) + log_ll = torch.empty(B, **f32) + 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, + token_logits, + targets, + source_lengths, + target_lengths, + durations, + blank_lp, + label_lp, + dur_lp, + token_lse, + alphas, + log_ll, ) ctx.blank_id = blank_id + ctx.duration_shape = duration_logits.shape + ctx.duration_dtype = duration_logits.dtype 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, + def backward(ctx, grad_loss: torch.Tensor): + ( + token_logits, + targets, + source_lengths, + target_lengths, + durations, + blank_lp, + label_lp, + dur_lp, + token_lse, + alphas, + log_ll, + ) = ctx.saved_tensors + if not (ctx.needs_input_grad[0] or ctx.needs_input_grad[1]): + return (None,) * 8 + + betas = torch.empty_like(alphas) + log_ll_bwd = torch.empty_like(log_ll) + ops.tdt_loss_bwd(blank_lp, label_lp, dur_lp, source_lengths, target_lengths, durations, betas, log_ll_bwd) + + grad_token_logits = torch.empty(token_logits.shape, device=token_logits.device, dtype=token_logits.dtype) + grad_duration_logits = torch.empty(ctx.duration_shape, device=token_logits.device, dtype=ctx.duration_dtype) + ops.tdt_logits_grad( + token_logits, + targets, + source_lengths, + target_lengths, + durations, + blank_lp, + label_lp, + dur_lp, + token_lse, + alphas, + betas, + log_ll, + grad_loss.float().contiguous(), + ctx.blank_id, + grad_token_logits, + grad_duration_logits, ) - - 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 @@ -108,42 +130,61 @@ def tdt_loss( targets: torch.Tensor, source_lengths: torch.Tensor, target_lengths: torch.Tensor, - durations: Union[List[int], torch.Tensor], + durations: Union[Sequence[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. + """Compute the TDT (Token-and-Duration Transducer) loss (https://arxiv.org/abs/2304.06795). 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]). + token_logits: Token logits of shape `(batch, T, U+1, vocab_size+1)`, in float32, float16 or bfloat16. + Only the last dimension needs to be contiguous. + duration_logits: Duration logits of shape `(batch, T, U+1, num_durations)`, same dtype as `token_logits`. + targets: Target labels of shape `(batch, U)`. + source_lengths: Encoder output lengths of shape `(batch,)`. + target_lengths: Target lengths of shape `(batch,)`. + durations: 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". + sigma: Logit undernormalization constant (see TDT paper). Defaults to `0.0`. + reduction: One of `"mean_volume"`, `"mean_batch"`, `"mean"`, `"sum"` or `"none"`, mirroring NeMo's + `RNNTLoss`. `"mean"` divides each loss by its target length before averaging over the batch. Returns: - Scalar loss tensor (or per-example losses if reduction="none"). + Scalar loss tensor (or per-example losses of shape `(batch,)` if `reduction="none"`). Samples without a + valid alignment get an infinite loss and a zero gradient. """ - if reduction not in ("mean", "sum", "none"): - raise ValueError(f'Invalid reduction mode "{reduction}". Expected one of "mean", "sum", or "none".') + if reduction not in _REDUCTIONS: + raise ValueError( + f'Invalid reduction mode "{reduction}". Expected one of {", ".join(repr(r) for r in _REDUCTIONS)}.' + ) - if isinstance(durations, (list, tuple)): - durations = torch.tensor(durations, dtype=torch.int32, device=token_logits.device) + device = token_logits.device + if isinstance(durations, torch.Tensor): + durations = durations.to(device=device, dtype=torch.int32).contiguous() 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, + durations = torch.tensor(list(durations), device=device, dtype=torch.int32) + targets = targets.to(device=device, dtype=torch.int32).contiguous() + source_lengths = source_lengths.to(device=device, dtype=torch.int32).contiguous() + target_lengths = target_lengths.to(device=device, dtype=torch.int32).contiguous() + + losses = TDTLoss.apply( + token_logits, + duration_logits, + targets, + source_lengths, + target_lengths, + durations, + int(blank_id), + float(sigma), ) + + if reduction == "mean_volume": + return losses.sum() / target_lengths.sum().float() + if reduction == "mean_batch": + return losses.mean() if reduction == "mean": - return (per_sample_loss / target_lengths.float()).mean() - elif reduction == "sum": - return per_sample_loss.sum() - return per_sample_loss + return (losses / target_lengths.float()).mean() + if reduction == "sum": + return losses.sum() + return losses diff --git a/tdt-loss/torch-ext/torch_binding.cpp b/tdt-loss/torch-ext/torch_binding.cpp index 5b7ca2ee1..f19d41276 100644 --- a/tdt-loss/torch-ext/torch_binding.cpp +++ b/tdt-loss/torch-ext/torch_binding.cpp @@ -7,8 +7,8 @@ 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) -> ()"); + "int blank_id, float sigma, Tensor! blank_lp, Tensor! label_lp, " + "Tensor! dur_lp, Tensor! token_lse) -> ()"); ops.impl("tdt_logprobs_fwd", torch::kCUDA, &tdt_logprobs_fwd); ops.def( @@ -20,23 +20,17 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { 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) -> ()"); + "Tensor! betas, Tensor! log_ll) -> ()"); 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, " + "tdt_logits_grad(Tensor token_logits, Tensor targets, " "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); + "Tensor blank_lp, Tensor label_lp, Tensor dur_lp, Tensor token_lse, " + "Tensor alphas, Tensor betas, Tensor log_ll, Tensor grad_loss, " + "int blank_id, Tensor! grad_token_logits, " + "Tensor! grad_duration_logits) -> ()"); + ops.impl("tdt_logits_grad", torch::kCUDA, &tdt_logits_grad); } REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/tdt-loss/torch-ext/torch_binding.h b/tdt-loss/torch-ext/torch_binding.h index 1e5b68f6b..d9c7d32d1 100644 --- a/tdt-loss/torch-ext/torch_binding.h +++ b/tdt-loss/torch-ext/torch_binding.h @@ -2,52 +2,42 @@ #include -// Forward: compute log-probabilities for blank, label, and duration +// Log-softmax gather over the vocabulary and duration heads. Writes the blank +// and label log-probabilities of every lattice node, the duration +// log-probabilities and the token log-normalizer (reused by the backward). 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); + torch::Tensor &label_lp, torch::Tensor &dur_lp, + torch::Tensor &token_lse); -// Forward pass of TDT loss (alpha computation) -void tdt_loss_fwd(torch::Tensor const &blank_lp, - torch::Tensor const &label_lp, +// Forward recursion over the lattice (alphas) and per-sample log-likelihood. +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, +// Backward recursion over the lattice (betas) and per-sample log-likelihood. +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); + torch::Tensor &log_ll); -// 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); +// Gradient of the per-sample losses (scaled by grad_loss) w.r.t. the token +// and duration logits. +void tdt_logits_grad( + torch::Tensor const &token_logits, torch::Tensor const &targets, + torch::Tensor const &source_lengths, torch::Tensor const &target_lengths, + torch::Tensor const &durations, torch::Tensor const &blank_lp, + torch::Tensor const &label_lp, torch::Tensor const &dur_lp, + torch::Tensor const &token_lse, torch::Tensor const &alphas, + torch::Tensor const &betas, torch::Tensor const &log_ll, + torch::Tensor const &grad_loss, int64_t blank_id, + torch::Tensor &grad_token_logits, torch::Tensor &grad_duration_logits); From 0bee960aec384bd0cb3ea3b30e89c7cba397fd39 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 28 Sep 2026 14:33:07 +1000 Subject: [PATCH 03/17] tdt-loss: add benchmark Co-Authored-By: Claude Opus 5.5 --- tdt-loss/benchmarks/benchmark.py | 145 +++++++++++++++++++++++++++++++ 1 file changed, 145 insertions(+) create mode 100644 tdt-loss/benchmarks/benchmark.py diff --git a/tdt-loss/benchmarks/benchmark.py b/tdt-loss/benchmarks/benchmark.py new file mode 100644 index 000000000..8d28cfbd1 --- /dev/null +++ b/tdt-loss/benchmarks/benchmark.py @@ -0,0 +1,145 @@ +import torch + +from kernels.benchmark import Benchmark + +DURATIONS = [0, 1, 2, 3, 4] + + +def tdt_loss_reference( + token_logits, duration_logits, targets, logit_lengths, target_lengths, blank_id, durations, sigma=0.0 +): + """Per-sample TDT loss, as implemented in `transformers.loss.loss_tdt.tdt_loss`.""" + device = token_logits.device + batch_size, max_t, max_u, _ = token_logits.shape + + token_log_probs = torch.log_softmax(token_logits.float(), dim=-1) - sigma + duration_log_probs = torch.log_softmax(duration_logits.float(), dim=-1) + + log_alpha = torch.full((batch_size, max_t, max_u), float("-inf"), device=device) + log_alpha[:, 0, 0] = 0.0 + blank_log_probs = token_log_probs[:, :, :, blank_id] + if max_u > 1: + targets_expanded = targets.long().unsqueeze(1).expand(-1, max_t, -1) + label_log_probs = torch.gather( + token_log_probs[:, :, : max_u - 1, :], dim=3, index=targets_expanded.unsqueeze(-1) + ).squeeze(-1) + + neg_inf = torch.tensor(float("-inf"), device=device) + for n in range(1, max_t + max_u - 1): + u_indices = torch.arange(max(0, n - max_t + 1), min(n + 1, max_u), device=device) + t_indices = n - u_indices + candidates = [] + for i, dur in enumerate(durations): + t_prev = t_indices - dur + valid_t = t_prev >= 0 + if not valid_t.any(): + continue + t_src = t_prev.clamp(min=0) + if dur > 0: + contrib = ( + log_alpha[:, t_src, u_indices] + + blank_log_probs[:, t_src, u_indices] + + duration_log_probs[:, t_src, u_indices, i] + ) + candidates.append(torch.where(valid_t.unsqueeze(0), contrib, neg_inf)) + valid_both = valid_t & (u_indices > 0) + if valid_both.any(): + u_src = (u_indices - 1).clamp(min=0) + u_src_label = u_src.clamp(max=max_u - 2) if max_u > 1 else u_src + contrib = ( + log_alpha[:, t_src, u_src] + + label_log_probs[:, t_src, u_src_label] + + duration_log_probs[:, t_src, u_src, i] + ) + candidates.append(torch.where(valid_both.unsqueeze(0), contrib, neg_inf)) + if candidates: + log_alpha[:, t_indices, u_indices] = torch.logsumexp(torch.stack(candidates), dim=0) + + batch_idx = torch.arange(batch_size, device=device) + target_lengths = target_lengths.long() + log_probs = torch.full((batch_size,), float("-inf"), device=device) + for i, dur in enumerate(durations): + if dur == 0: + continue + t_final = logit_lengths.long() - dur + valid = t_final >= 0 + if not valid.any(): + continue + t_clamped = t_final.clamp(min=0) + terminal = ( + log_alpha[batch_idx, t_clamped, target_lengths] + + token_log_probs[batch_idx, t_clamped, target_lengths, blank_id] + + duration_log_probs[batch_idx, t_clamped, target_lengths, i] + ) + log_probs = torch.where(valid, torch.logsumexp(torch.stack([log_probs, terminal]), dim=0), log_probs) + return -log_probs + + +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, + DURATIONS, + 0, + 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() From 34479454aa9c5f349ad5cc7dd3185c772aa80f1e Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 28 Sep 2026 14:36:03 +1000 Subject: [PATCH 04/17] tdt-loss: fix test references for unreachable nodes and long targets Co-Authored-By: Claude Opus 5.5 --- tdt-loss/tests/test_tdt_loss.py | 57 ++++++++++++++++++++++----------- 1 file changed, 38 insertions(+), 19 deletions(-) diff --git a/tdt-loss/tests/test_tdt_loss.py b/tdt-loss/tests/test_tdt_loss.py index 580fbffa6..b08baf26a 100644 --- a/tdt-loss/tests/test_tdt_loss.py +++ b/tdt-loss/tests/test_tdt_loss.py @@ -24,6 +24,8 @@ DEVICE = torch.device("cuda") DURATIONS = [0, 1, 2, 3, 4] REDUCTIONS = ["mean_volume", "mean_batch", "mean", "sum", "none"] +# float32 lattice sums over ~1000 labels lose a few digits. +GRAD_ATOL_LONG = 1e-2 def reduce(losses, target_lengths, reduction): @@ -40,16 +42,24 @@ def reduce(losses, target_lengths, reduction): def tdt_loss_reference( - token_logits, duration_logits, targets, logit_lengths, target_lengths, blank_id, durations, sigma=0.0 + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + blank_id, + durations, + sigma=0.0, + dtype=torch.float32, ): - """Per-sample TDT loss, as implemented in `transformers.loss.loss_tdt.tdt_loss`.""" + """Per-sample TDT loss, as implemented in `transformers.loss.loss_tdt.tdt_loss` (computed in `dtype`).""" device = token_logits.device batch_size, max_t, max_u, _ = token_logits.shape - token_log_probs = torch.log_softmax(token_logits.float(), dim=-1) - sigma - duration_log_probs = torch.log_softmax(duration_logits.float(), dim=-1) + token_log_probs = torch.log_softmax(token_logits.to(dtype), dim=-1) - sigma + duration_log_probs = torch.log_softmax(duration_logits.to(dtype), dim=-1) - log_alpha = torch.full((batch_size, max_t, max_u), float("-inf"), device=device) + log_alpha = torch.full((batch_size, max_t, max_u), float("-inf"), device=device, dtype=dtype) log_alpha[:, 0, 0] = 0.0 blank_log_probs = token_log_probs[:, :, :, blank_id] if max_u > 1: @@ -58,7 +68,7 @@ def tdt_loss_reference( token_log_probs[:, :, : max_u - 1, :], dim=3, index=targets_expanded.unsqueeze(-1) ).squeeze(-1) - neg_inf = torch.tensor(float("-inf"), device=device) + neg_inf = torch.tensor(float("-inf"), device=device, dtype=dtype) for n in range(1, max_t + max_u - 1): u_indices = torch.arange(max(0, n - max_t + 1), min(n + 1, max_u), device=device) t_indices = n - u_indices @@ -91,7 +101,7 @@ def tdt_loss_reference( batch_idx = torch.arange(batch_size, device=device) target_lengths = target_lengths.long() - log_probs = torch.full((batch_size,), float("-inf"), device=device) + log_probs = torch.full((batch_size,), float("-inf"), device=device, dtype=dtype) for i, dur in enumerate(durations): if dur == 0: continue @@ -148,6 +158,12 @@ def log_add(a, b): return torch.tensor(losses, dtype=torch.float64) +def assert_grad_close(actual, expected, **kwargs): + # The reference gets NaN gradients on lattice nodes that cannot be reached (e.g. u > t when there is no + # zero duration), from the backward of a logsumexp over only -inf. The true gradient there is zero. + torch.testing.assert_close(actual.to(expected.dtype), expected.nan_to_num(nan=0.0), **kwargs) + + def make_inputs(B, T, U, V, durations, dtype=torch.float32, seed=0, joint=False): """Random inputs with variable lengths. `joint=True` returns slices of a single joint tensor, like `ParakeetForTDT` does.""" @@ -210,8 +226,8 @@ def test_forward_backward_match_reference(durations, sigma, reduction): loss.sum().backward() torch.testing.assert_close(loss, expected, atol=1e-4, rtol=1e-4) - torch.testing.assert_close(tok.grad, ref_tok.grad, atol=1e-5, rtol=1e-4) - torch.testing.assert_close(dur.grad, ref_dur.grad, atol=1e-5, rtol=1e-4) + assert_grad_close(tok.grad, ref_tok.grad, atol=1e-5, rtol=1e-4) + assert_grad_close(dur.grad, ref_dur.grad, atol=1e-5, rtol=1e-4) def test_joint_logits_slices(): @@ -239,7 +255,7 @@ def test_joint_logits_slices(): expected.backward() torch.testing.assert_close(loss, expected, atol=1e-4, rtol=1e-4) - torch.testing.assert_close(joint.grad, ref_joint.grad, atol=1e-5, rtol=1e-4) + assert_grad_close(joint.grad, ref_joint.grad, atol=1e-5, rtol=1e-4) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -268,8 +284,8 @@ def test_half_precision(dtype): expected.backward() torch.testing.assert_close(loss, expected, atol=1e-4, rtol=1e-4) - torch.testing.assert_close(tok.grad.float(), ref_tok.grad, atol=1e-3, rtol=1e-2) - torch.testing.assert_close(dur.grad.float(), ref_dur.grad, atol=1e-3, rtol=1e-2) + assert_grad_close(tok.grad, ref_tok.grad, atol=1e-3, rtol=1e-2) + assert_grad_close(dur.grad, ref_dur.grad, atol=1e-3, rtol=1e-2) def test_parakeet_like_shapes(): @@ -288,8 +304,8 @@ def test_parakeet_like_shapes(): expected.sum().backward() torch.testing.assert_close(losses, expected, atol=1e-3, rtol=1e-4) - torch.testing.assert_close(tok.grad, ref_tok.grad, atol=1e-5, rtol=1e-4) - torch.testing.assert_close(dur.grad, ref_dur.grad, atol=1e-5, rtol=1e-4) + assert_grad_close(tok.grad, ref_tok.grad, atol=1e-5, rtol=1e-4) + assert_grad_close(dur.grad, ref_dur.grad, atol=1e-5, rtol=1e-4) def test_long_targets(): @@ -300,12 +316,15 @@ def test_long_targets(): loss = run_kernel(tok, duration_logits, targets, logit_lengths, target_lengths, DURATIONS, 0.0, "sum") loss.backward() - ref_tok = token_logits.clone().requires_grad_(True) - expected = tdt_loss_reference(ref_tok, duration_logits, targets, logit_lengths, target_lengths, 0, DURATIONS) - expected.sum().backward() + # The log-likelihood is in the thousands, so compare against a float64 reference. + ref_tok = token_logits.double().requires_grad_(True) + expected = tdt_loss_reference( + ref_tok, duration_logits, targets, logit_lengths, target_lengths, 0, DURATIONS, dtype=torch.float64 + ).sum() + expected.backward() - torch.testing.assert_close(loss, expected.sum(), atol=1e-2, rtol=1e-4) - torch.testing.assert_close(tok.grad, ref_tok.grad, atol=1e-5, rtol=1e-4) + torch.testing.assert_close(loss.double(), expected, atol=1e-2, rtol=1e-5) + assert_grad_close(tok.grad, ref_tok.grad, atol=GRAD_ATOL_LONG, rtol=0) def test_empty_targets(): From 8e5cb43a95dafd0ab4101f2e64b1a5d534807241 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 28 Sep 2026 14:39:34 +1000 Subject: [PATCH 05/17] tdt-loss: run the lattice recursions in float64 alpha + beta - log_likelihood cancels in float32 when the log-likelihood is large (long targets). The recursions are ~2% of the runtime, so accumulate them in float64; log-probs and logits stay in float32 or lower. Co-Authored-By: Claude Opus 5.5 --- tdt-loss/tdt_loss/common.cuh | 5 ++- tdt-loss/tdt_loss/lattice.cu | 41 +++++++++++++--------- tdt-loss/tdt_loss/logprobs.cu | 46 +++++++++++++++---------- tdt-loss/tests/test_tdt_loss.py | 4 +-- tdt-loss/torch-ext/tdt_loss/__init__.py | 8 +++-- 5 files changed, 62 insertions(+), 42 deletions(-) diff --git a/tdt-loss/tdt_loss/common.cuh b/tdt-loss/tdt_loss/common.cuh index e593225d4..208b72fda 100644 --- a/tdt-loss/tdt_loss/common.cuh +++ b/tdt-loss/tdt_loss/common.cuh @@ -14,11 +14,10 @@ constexpr int kMaxDurations = 64; __device__ __forceinline__ float neg_inf() { return -INFINITY; } // log(exp(a) + exp(b)), safe when either argument is -inf. -__device__ __forceinline__ float log_add(float a, float b) { +__device__ __forceinline__ double log_add(double a, double b) { if (a == -INFINITY) return b; if (b == -INFINITY) return a; - const float m = fmaxf(a, b); - return m + log1pf(expf(-fabsf(a - b))); + return fmax(a, b) + log1p(exp(-fabs(a - b))); } // Merge two partial (max, sum(exp(x - max))) pairs of an online softmax. diff --git a/tdt-loss/tdt_loss/lattice.cu b/tdt-loss/tdt_loss/lattice.cu index 23de82232..2329179d8 100644 --- a/tdt-loss/tdt_loss/lattice.cu +++ b/tdt-loss/tdt_loss/lattice.cu @@ -9,6 +9,10 @@ // Every arc strictly increases t + u, so all nodes of an anti-diagonal // t + u = n are independent. One thread block handles one sample and sweeps // the anti-diagonals, with the threads striding over u. +// +// The recursions accumulate in float64: they are a small fraction of the total +// cost, and the gradient combines alpha + beta - log_likelihood, which cancels +// catastrophically in float32 when the log-likelihood is large (long targets). #include #include @@ -34,8 +38,8 @@ __global__ void tdt_alpha_kernel(const float *__restrict__ blank_lp, const int *__restrict__ source_lengths, const int *__restrict__ target_lengths, const int *__restrict__ durations, int max_T, - int max_U, int D, float *__restrict__ alphas, - float *__restrict__ log_ll) { + int max_U, int D, double *__restrict__ alphas, + double *__restrict__ log_ll) { __shared__ int shm_durations[kMaxDurations]; const int b = blockIdx.x; @@ -47,14 +51,14 @@ __global__ void tdt_alpha_kernel(const float *__restrict__ blank_lp, const float *blank = blank_lp + sample; const float *label = label_lp + sample; const float *dur = dur_lp + sample * D; - float *alpha = alphas + sample; + double *alpha = alphas + sample; for (int n = 0; n < T + U; ++n) { const int u_min = max(0, n - T + 1); const int u_max = min(n, U); for (int u = u_min + threadIdx.x; u <= u_max; u += blockDim.x) { const int t = n - u; - float a = n == 0 ? 0.f : neg_inf(); + double a = n == 0 ? 0.0 : -INFINITY; for (int i = 0; i < D && n > 0; ++i) { const int d = shm_durations[i]; const int t_prev = t - d; @@ -74,7 +78,7 @@ __global__ void tdt_alpha_kernel(const float *__restrict__ blank_lp, } if (threadIdx.x == 0) { - float ll = neg_inf(); + double ll = -INFINITY; for (int i = 0; i < D; ++i) { const int d = shm_durations[i]; const int t_prev = T - d; @@ -92,8 +96,8 @@ __global__ void tdt_beta_kernel(const float *__restrict__ blank_lp, const int *__restrict__ source_lengths, const int *__restrict__ target_lengths, const int *__restrict__ durations, int max_T, - int max_U, int D, float *__restrict__ betas, - float *__restrict__ log_ll) { + int max_U, int D, double *__restrict__ betas, + double *__restrict__ log_ll) { __shared__ int shm_durations[kMaxDurations]; const int b = blockIdx.x; @@ -105,7 +109,7 @@ __global__ void tdt_beta_kernel(const float *__restrict__ blank_lp, const float *blank = blank_lp + sample; const float *label = label_lp + sample; const float *dur = dur_lp + sample * D; - float *beta = betas + sample; + double *beta = betas + sample; for (int n = T + U - 1; n >= 0; --n) { const int u_min = max(0, n - T + 1); @@ -113,17 +117,17 @@ __global__ void tdt_beta_kernel(const float *__restrict__ blank_lp, for (int u = u_min + threadIdx.x; u <= u_max; u += blockDim.x) { const int t = n - u; const int64_t node = static_cast(t) * max_U + u; - float bt = neg_inf(); + double bt = -INFINITY; for (int i = 0; i < D; ++i) { const int d = shm_durations[i]; const int t_next = t + d; - const float lp_dur = dur[node * D + i]; + const double lp_dur = dur[node * D + i]; if (d > 0) { if (t_next < T) { const int64_t dst = static_cast(t_next) * max_U + u; bt = log_add(bt, beta[dst] + blank[node] + lp_dur); } else if (t_next == T && u == U) { - bt = log_add(bt, blank[node] + lp_dur); + bt = log_add(bt, static_cast(blank[node]) + lp_dur); } } if (u < U && t_next < T) { @@ -137,7 +141,7 @@ __global__ void tdt_beta_kernel(const float *__restrict__ blank_lp, } if (threadIdx.x == 0) { - log_ll[b] = T > 0 ? beta[0] : neg_inf(); + log_ll[b] = T > 0 ? beta[0] : -INFINITY; } } @@ -159,10 +163,15 @@ void check_lattice_args(torch::Tensor const &blank_lp, for (auto const *x : {&blank_lp, &label_lp, &dur_lp, &out, &log_ll}) { TORCH_CHECK(x->device() == blank_lp.device(), "all tensors must be on the same device"); - TORCH_CHECK(x->scalar_type() == torch::kFloat32, - "lattice tensors must be float32"); TORCH_CHECK(x->is_contiguous(), "lattice tensors must be contiguous"); } + for (auto const *x : {&blank_lp, &label_lp, &dur_lp}) { + TORCH_CHECK(x->scalar_type() == torch::kFloat32, + "log-prob tensors must be float32"); + } + TORCH_CHECK(out.scalar_type() == torch::kFloat64 && + log_ll.scalar_type() == torch::kFloat64, + "alphas/betas and log_ll must be float64"); for (auto const *x : {&source_lengths, &target_lengths, &durations}) { TORCH_CHECK(x->device() == blank_lp.device(), "all tensors must be on the same device"); @@ -206,7 +215,7 @@ void tdt_loss_fwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, dur_lp.data_ptr(), source_lengths.data_ptr(), target_lengths.data_ptr(), durations.data_ptr(), blank_lp.size(1), blank_lp.size(2), durations.numel(), - alphas.data_ptr(), log_ll.data_ptr()); + alphas.data_ptr(), log_ll.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -229,6 +238,6 @@ void tdt_loss_bwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, dur_lp.data_ptr(), source_lengths.data_ptr(), target_lengths.data_ptr(), durations.data_ptr(), blank_lp.size(1), blank_lp.size(2), durations.numel(), - betas.data_ptr(), log_ll.data_ptr()); + betas.data_ptr(), log_ll.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/tdt-loss/tdt_loss/logprobs.cu b/tdt-loss/tdt_loss/logprobs.cu index 2f93202c6..9dce7a7d9 100644 --- a/tdt-loss/tdt_loss/logprobs.cu +++ b/tdt-loss/tdt_loss/logprobs.cu @@ -104,8 +104,8 @@ __global__ void tdt_logits_grad_kernel( const int *__restrict__ target_lengths, const int *__restrict__ durations, const float *__restrict__ blank_lp, const float *__restrict__ label_lp, const float *__restrict__ dur_lp, const float *__restrict__ token_lse, - const float *__restrict__ alphas, const float *__restrict__ betas, - const float *__restrict__ log_ll, const float *__restrict__ grad_loss, + const double *__restrict__ alphas, const double *__restrict__ betas, + const double *__restrict__ log_ll, const float *__restrict__ grad_loss, int max_T, int max_U, int V, int D, int blank_id, scalar_t *__restrict__ grad_token_logits, scalar_t *__restrict__ grad_duration_logits) { @@ -126,7 +126,7 @@ __global__ void tdt_logits_grad_kernel( scalar_t *gd = grad_duration_logits + row * D; const scalar_t zero = static_cast(0.f); - const float ll = log_ll[b]; + const double ll = log_ll[b]; // Infeasible samples (infinite loss) get a zero gradient instead of NaNs. if (t >= T || u > U || !isfinite(ll)) { for (int v = threadIdx.x; v < V; v += blockDim.x) gx[v] = zero; @@ -135,33 +135,33 @@ __global__ void tdt_logits_grad_kernel( } const int64_t sample = static_cast(b) * max_T * max_U; - const float alpha = alphas[row]; - const float blank = blank_lp[row]; - const float label = label_lp[row]; + const double alpha = alphas[row]; + const double blank = blank_lp[row]; + const double label = label_lp[row]; for (int i = threadIdx.x; i < D; i += blockDim.x) { const int d = durations[i]; const int t_next = t + d; - const float lp_dur = dur_lp[row * D + i]; + const double lp_dur = dur_lp[row * D + i]; float p_blank = 0.f; if (d > 0) { - float beta_next = neg_inf(); + double beta_next = -INFINITY; if (t_next < T) { beta_next = betas[sample + static_cast(t_next) * max_U + u]; } else if (t_next == T && u == U) { - beta_next = 0.f; // terminal arc + beta_next = 0.0; // terminal arc } if (beta_next != -INFINITY) { - p_blank = expf(alpha + blank + lp_dur + beta_next - ll); + p_blank = static_cast(exp(alpha + blank + lp_dur + beta_next - ll)); } } float p_label = 0.f; if (u < U && t_next < T) { - const float beta_next = + const double beta_next = betas[sample + static_cast(t_next) * max_U + u + 1]; - p_label = expf(alpha + label + lp_dur + beta_next - ll); + p_label = static_cast(exp(alpha + label + lp_dur + beta_next - ll)); } shm_blank[i] = p_blank; @@ -341,13 +341,23 @@ void tdt_logits_grad( check_int_tensor(source_lengths, "source_lengths", token_logits); check_int_tensor(target_lengths, "target_lengths", token_logits); check_int_tensor(durations, "durations", token_logits); - for (auto *in : {&blank_lp, &label_lp, &token_lse, &alphas, &betas}) { - check_float_out(*in, "lattice tensor", token_logits); + for (auto *in : {&blank_lp, &label_lp, &token_lse}) { + check_float_out(*in, "log-prob tensor", token_logits); TORCH_CHECK(in->numel() == B * max_T * max_U, - "lattice tensors must have shape (batch, T, U+1)"); + "log-prob tensors must have shape (batch, T, U+1)"); + } + for (auto *in : {&alphas, &betas}) { + TORCH_CHECK(in->device() == token_logits.device() && + in->scalar_type() == torch::kFloat64 && + in->is_contiguous() && in->numel() == B * max_T * max_U, + "alphas and betas must be contiguous float64 tensors of shape " + "(batch, T, U+1)"); } check_float_out(dur_lp, "dur_lp", token_logits); - check_float_out(log_ll, "log_ll", token_logits); + TORCH_CHECK(log_ll.device() == token_logits.device() && + log_ll.scalar_type() == torch::kFloat64 && + log_ll.is_contiguous(), + "log_ll must be a contiguous float64 tensor"); check_float_out(grad_loss, "grad_loss", token_logits); TORCH_CHECK(log_ll.numel() == B && grad_loss.numel() == B, "log_ll and grad_loss must have shape (batch,)"); @@ -381,8 +391,8 @@ void tdt_logits_grad( source_lengths.data_ptr(), target_lengths.data_ptr(), durations.data_ptr(), blank_lp.data_ptr(), label_lp.data_ptr(), dur_lp.data_ptr(), - token_lse.data_ptr(), alphas.data_ptr(), - betas.data_ptr(), log_ll.data_ptr(), + token_lse.data_ptr(), alphas.data_ptr(), + betas.data_ptr(), log_ll.data_ptr(), grad_loss.data_ptr(), max_T, max_U, V, D, blank_id, grad_token_logits.data_ptr(), grad_duration_logits.data_ptr()); diff --git a/tdt-loss/tests/test_tdt_loss.py b/tdt-loss/tests/test_tdt_loss.py index b08baf26a..39ff3d66b 100644 --- a/tdt-loss/tests/test_tdt_loss.py +++ b/tdt-loss/tests/test_tdt_loss.py @@ -24,8 +24,8 @@ DEVICE = torch.device("cuda") DURATIONS = [0, 1, 2, 3, 4] REDUCTIONS = ["mean_volume", "mean_batch", "mean", "sum", "none"] -# float32 lattice sums over ~1000 labels lose a few digits. -GRAD_ATOL_LONG = 1e-2 +# Gradient tolerance against a float64 reference for ~1000 labels (float32 log-probs, float64 lattice). +GRAD_ATOL_LONG = 1e-3 def reduce(losses, target_lengths, reduction): diff --git a/tdt-loss/torch-ext/tdt_loss/__init__.py b/tdt-loss/torch-ext/tdt_loss/__init__.py index 07b026a97..841bab359 100644 --- a/tdt-loss/torch-ext/tdt_loss/__init__.py +++ b/tdt-loss/torch-ext/tdt_loss/__init__.py @@ -57,8 +57,10 @@ def forward( token_lse, ) - alphas = torch.empty(B, max_T, max_U, **f32) - log_ll = torch.empty(B, **f32) + # The lattice recursions run in float64, see lattice.cu. + f64 = dict(device=token_logits.device, dtype=torch.float64) + alphas = torch.empty(B, max_T, max_U, **f64) + log_ll = torch.empty(B, **f64) ops.tdt_loss_fwd(blank_lp, label_lp, dur_lp, source_lengths, target_lengths, durations, alphas, log_ll) ctx.save_for_backward( @@ -77,7 +79,7 @@ def forward( ctx.blank_id = blank_id ctx.duration_shape = duration_logits.shape ctx.duration_dtype = duration_logits.dtype - return -log_ll + return -log_ll.float() @staticmethod def backward(ctx, grad_loss: torch.Tensor): From 3301b5d7cbcf726802fb572fc46c1a8ca6648626 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 28 Sep 2026 17:02:11 +1000 Subject: [PATCH 06/17] tdt-loss: keep float32 transcendentals in the lattice recursions Only the log-values are float64; exp/log act on differences in float32, and each node uses a two-pass logsumexp instead of chained log_add calls. Co-Authored-By: Claude Opus 5.5 --- tdt-loss/tdt_loss/common.cuh | 19 +++++++++-- tdt-loss/tdt_loss/lattice.cu | 64 +++++++++++++++++++---------------- tdt-loss/tdt_loss/logprobs.cu | 4 +-- 3 files changed, 53 insertions(+), 34 deletions(-) diff --git a/tdt-loss/tdt_loss/common.cuh b/tdt-loss/tdt_loss/common.cuh index 208b72fda..70d2d36ae 100644 --- a/tdt-loss/tdt_loss/common.cuh +++ b/tdt-loss/tdt_loss/common.cuh @@ -13,11 +13,26 @@ constexpr int kMaxDurations = 64; __device__ __forceinline__ float neg_inf() { return -INFINITY; } -// log(exp(a) + exp(b)), safe when either argument is -inf. +// log(exp(a) + exp(b)), safe when either argument is -inf. The arguments are +// float64 to keep large log-likelihoods exact, but the correction term only +// depends on |a - b| and is computed in float32. __device__ __forceinline__ double log_add(double a, double b) { if (a == -INFINITY) return b; if (b == -INFINITY) return a; - return fmax(a, b) + log1p(exp(-fabs(a - b))); + return fmax(a, b) + log1pf(expf(-fabsf(static_cast(a - b)))); +} + +// log(sum(exp(v))) over the values produced by `for_each(f)`, which calls f(v) +// for every value. Two passes (max, then sum) keep the exponentials +// independent of each other instead of chaining log_add calls. +template +__device__ __forceinline__ double log_sum_exp(ForEach for_each) { + double m = -INFINITY; + for_each([&](double v) { m = fmax(m, v); }); + if (m == -INFINITY) return m; + float s = 0.f; + for_each([&](double v) { s += expf(static_cast(v - m)); }); + return m + logf(s); } // Merge two partial (max, sum(exp(x - max))) pairs of an online softmax. diff --git a/tdt-loss/tdt_loss/lattice.cu b/tdt-loss/tdt_loss/lattice.cu index 2329179d8..be3eed1fa 100644 --- a/tdt-loss/tdt_loss/lattice.cu +++ b/tdt-loss/tdt_loss/lattice.cu @@ -58,20 +58,23 @@ __global__ void tdt_alpha_kernel(const float *__restrict__ blank_lp, const int u_max = min(n, U); for (int u = u_min + threadIdx.x; u <= u_max; u += blockDim.x) { const int t = n - u; - double a = n == 0 ? 0.0 : -INFINITY; - for (int i = 0; i < D && n > 0; ++i) { - const int d = shm_durations[i]; - const int t_prev = t - d; - if (t_prev < 0) continue; - if (d > 0) { - const int64_t src = static_cast(t_prev) * max_U + u; - a = log_add(a, alpha[src] + blank[src] + dur[src * D + i]); - } - if (u > 0) { - const int64_t src = static_cast(t_prev) * max_U + u - 1; - a = log_add(a, alpha[src] + label[src] + dur[src * D + i]); - } - } + const double a = + n == 0 ? 0.0 : log_sum_exp([&](auto &&f) { + for (int i = 0; i < D; ++i) { + const int d = shm_durations[i]; + const int t_prev = t - d; + if (t_prev < 0) continue; + if (d > 0) { + const int64_t src = static_cast(t_prev) * max_U + u; + f(alpha[src] + blank[src] + dur[src * D + i]); + } + if (u > 0) { + const int64_t src = + static_cast(t_prev) * max_U + u - 1; + f(alpha[src] + label[src] + dur[src * D + i]); + } + } + }); alpha[static_cast(t) * max_U + u] = a; } __syncthreads(); @@ -117,24 +120,25 @@ __global__ void tdt_beta_kernel(const float *__restrict__ blank_lp, for (int u = u_min + threadIdx.x; u <= u_max; u += blockDim.x) { const int t = n - u; const int64_t node = static_cast(t) * max_U + u; - double bt = -INFINITY; - for (int i = 0; i < D; ++i) { - const int d = shm_durations[i]; - const int t_next = t + d; - const double lp_dur = dur[node * D + i]; - if (d > 0) { - if (t_next < T) { - const int64_t dst = static_cast(t_next) * max_U + u; - bt = log_add(bt, beta[dst] + blank[node] + lp_dur); - } else if (t_next == T && u == U) { - bt = log_add(bt, static_cast(blank[node]) + lp_dur); + const double bt = log_sum_exp([&](auto &&f) { + for (int i = 0; i < D; ++i) { + const int d = shm_durations[i]; + const int t_next = t + d; + const double lp_dur = dur[node * D + i]; + if (d > 0) { + if (t_next < T) { + const int64_t dst = static_cast(t_next) * max_U + u; + f(beta[dst] + blank[node] + lp_dur); + } else if (t_next == T && u == U) { + f(blank[node] + lp_dur); // terminal arc + } + } + if (u < U && t_next < T) { + const int64_t dst = static_cast(t_next) * max_U + u + 1; + f(beta[dst] + label[node] + lp_dur); } } - if (u < U && t_next < T) { - const int64_t dst = static_cast(t_next) * max_U + u + 1; - bt = log_add(bt, beta[dst] + label[node] + lp_dur); - } - } + }); beta[node] = bt; } __syncthreads(); diff --git a/tdt-loss/tdt_loss/logprobs.cu b/tdt-loss/tdt_loss/logprobs.cu index 9dce7a7d9..ec25f996a 100644 --- a/tdt-loss/tdt_loss/logprobs.cu +++ b/tdt-loss/tdt_loss/logprobs.cu @@ -153,7 +153,7 @@ __global__ void tdt_logits_grad_kernel( beta_next = 0.0; // terminal arc } if (beta_next != -INFINITY) { - p_blank = static_cast(exp(alpha + blank + lp_dur + beta_next - ll)); + p_blank = expf(static_cast(alpha + blank + lp_dur + beta_next - ll)); } } @@ -161,7 +161,7 @@ __global__ void tdt_logits_grad_kernel( if (u < U && t_next < T) { const double beta_next = betas[sample + static_cast(t_next) * max_U + u + 1]; - p_label = static_cast(exp(alpha + label + lp_dur + beta_next - ll)); + p_label = expf(static_cast(alpha + label + lp_dur + beta_next - ll)); } shm_blank[i] = p_blank; From da05018fbe9469ff5cbb6019bdc097d07cdc9cf5 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 28 Sep 2026 17:10:56 +1000 Subject: [PATCH 07/17] tdt-loss: update README Co-Authored-By: Claude Opus 5.5 --- tdt-loss/README.md | 73 +++++++++++++++++++++++++++++++--------------- 1 file changed, 49 insertions(+), 24 deletions(-) diff --git a/tdt-loss/README.md b/tdt-loss/README.md index 522601157..02fb59782 100644 --- a/tdt-loss/README.md +++ b/tdt-loss/README.md @@ -2,7 +2,17 @@ tags: - kernel --- -This CUDA extension implements TDT (Token-and-Duration Transducer) loss, originally from [eustlb/tdt-loss](https://huggingface.co/eustlb/tdt-loss). +CUDA kernels for the TDT (Token-and-Duration Transducer, [arXiv:2304.06795](https://arxiv.org/abs/2304.06795)) loss, +used to train [Parakeet TDT](https://huggingface.co/docs/transformers/model_doc/parakeet) models in 🤗 Transformers. + +The forward and backward passes are computed with dedicated kernels: + +- a log-softmax over the vocabulary for every lattice node (one thread block per node, online softmax), +- the alpha / beta recursions over the anti-diagonals of the lattice (one thread block per sample), +- a fused gradient w.r.t. the token and duration logits. + +Logits can be float32, float16 or bfloat16 and only need a contiguous last dimension, so slices of a joint +`(..., vocab_size + num_durations)` output are used without a copy. The lattice recursions run in float64. ## Usage @@ -16,32 +26,47 @@ This CUDA extension implements TDT (Token-and-Duration Transducer) loss, origina 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", +tdt = get_kernel("kernels-community/tdt-loss", version=1) + +batch_size, max_frames, max_labels, vocab_size = 2, 100, 20, 1024 +durations = [0, 1, 2, 3, 4] +blank_id = vocab_size # blank is the last token +device = "cuda" + +# Joint network output: vocabulary (with blank) followed by the durations. +logits = torch.randn( + batch_size, max_frames, max_labels + 1, vocab_size + 1 + len(durations), device=device, requires_grad=True +) +targets = torch.randint(0, vocab_size, (batch_size, max_labels), device=device) +logit_lengths = torch.tensor([max_frames, max_frames - 10], device=device) +target_lengths = torch.tensor([max_labels, max_labels - 5], device=device) + +loss = tdt.tdt_loss( + logits[..., : vocab_size + 1], + logits[..., vocab_size + 1 :], + targets, + logit_lengths, + target_lengths, + durations, + blank_id, + sigma=0.0, + reduction="mean", # or "mean_volume", "mean_batch", "sum", "none" ) loss.backward() ``` +Samples without a valid alignment get an infinite loss and a zero gradient. + ## Benchmarks -See the [benchmark script](https://gist.github.com/eustlb/42bb54363e82ad6c0f6fc8e24f604c18). +Forward + backward on an NVIDIA L4, with the logits given as slices of a joint output, compared with the vectorized +PyTorch implementation from `transformers.loss.loss_tdt`. Run `benchmarks/benchmark.py` with `kernels benchmark`. + +| batch, frames, labels, vocab | dtype | PyTorch | kernel | speedup | +|---|---|---:|---:|---:| +| 8, 100, 20, 1025 | float32 | 536 ms | 4.3 ms | 125x | +| 4, 200, 40, 8193 | float32 | 1252 ms | 53 ms | 24x | +| 4, 200, 40, 8193 | bfloat16 | 1242 ms | 39 ms | 32x | + +The kernels run close to the memory bandwidth of the GPU; most of the remaining time is spent by PyTorch +materializing the gradient of the two logits slices. From a76e9fa69dcc07abb9a1eaf1746a88ff74d96675 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 28 Sep 2026 17:14:47 +1000 Subject: [PATCH 08/17] tdt-loss: load the kernel with get_kernel in tests, add kernels_ci markers and flake.lock Co-Authored-By: Claude Opus 5.5 --- tdt-loss/flake.lock | 117 ++++++++++++++++++++++++++++++++ tdt-loss/tests/test_tdt_loss.py | 19 +++--- 2 files changed, 128 insertions(+), 8 deletions(-) create mode 100644 tdt-loss/flake.lock diff --git a/tdt-loss/flake.lock b/tdt-loss/flake.lock new file mode 100644 index 000000000..35a669620 --- /dev/null +++ b/tdt-loss/flake.lock @@ -0,0 +1,117 @@ +{ + "nodes": { + "flake-compat": { + "locked": { + "lastModified": 1767039857, + "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=", + "owner": "edolstra", + "repo": "flake-compat", + "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab", + "type": "github" + }, + "original": { + "owner": "edolstra", + "repo": "flake-compat", + "type": "github" + } + }, + "flake-utils": { + "inputs": { + "systems": "systems" + }, + "locked": { + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "kernel-builder": { + "inputs": { + "flake-compat": "flake-compat", + "flake-utils": "flake-utils", + "nixpkgs": "nixpkgs", + "rust-overlay": "rust-overlay" + }, + "locked": { + "lastModified": 1788544112, + "narHash": "sha256-aqwWRYCri6ZTSMQNRE0EQZRvP45nubvcW62pVhSeoyg=", + "owner": "huggingface", + "repo": "kernels", + "rev": "a7f0afdb29a6a3372b1d47180cc0c182454c5e3b", + "type": "github" + }, + "original": { + "owner": "huggingface", + "repo": "kernels", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1783284758, + "narHash": "sha256-tiQ8/qi8I45OOaBBYlVbXoAVkeQzvvTQOv5I45rMw5o=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d", + "type": "github" + }, + "original": { + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d", + "type": "github" + } + }, + "root": { + "inputs": { + "kernel-builder": "kernel-builder" + } + }, + "rust-overlay": { + "inputs": { + "nixpkgs": [ + "kernel-builder", + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1783320166, + "narHash": "sha256-l7C/OsjcnWDOk2K3ssj+SBduwL67LashjBqis9+t468=", + "owner": "oxalica", + "repo": "rust-overlay", + "rev": "20ee15370c9256669d66968b89ee20a4b0a4e673", + "type": "github" + }, + "original": { + "owner": "oxalica", + "repo": "rust-overlay", + "type": "github" + } + }, + "systems": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/tdt-loss/tests/test_tdt_loss.py b/tdt-loss/tests/test_tdt_loss.py index 39ff3d66b..8bf140267 100644 --- a/tdt-loss/tests/test_tdt_loss.py +++ b/tdt-loss/tests/test_tdt_loss.py @@ -6,20 +6,15 @@ - `tdt_loss_naive`: a loop-based implementation that is easy to verify by hand. """ -from pathlib import Path - import pytest import torch +import kernels + if not torch.cuda.is_available(): pytest.skip("CUDA not available", allow_module_level=True) -try: - import tdt_loss as tdt_loss_kernel -except ImportError: - from kernels import get_local_kernel - - tdt_loss_kernel = get_local_kernel(repo_path=Path(__file__).parent.parent, package_name="tdt_loss") +tdt_loss_kernel = kernels.get_kernel("kernels-community/tdt-loss", version=1) DEVICE = torch.device("cuda") DURATIONS = [0, 1, 2, 3, 4] @@ -194,6 +189,7 @@ def run_kernel(token_logits, duration_logits, targets, logit_lengths, target_len ) +@pytest.mark.kernels_ci @pytest.mark.parametrize("sigma", [0.0, 0.05]) def test_matches_naive(sigma): token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs(3, 12, 5, 16, DURATIONS) @@ -206,6 +202,7 @@ def test_matches_naive(sigma): torch.testing.assert_close(losses.double().cpu(), expected, atol=1e-4, rtol=1e-5) +@pytest.mark.kernels_ci @pytest.mark.parametrize("reduction", REDUCTIONS) @pytest.mark.parametrize("sigma", [0.0, 0.05]) @pytest.mark.parametrize("durations", [DURATIONS, [1, 2], [0, 1, 3, 5]]) @@ -230,6 +227,7 @@ def test_forward_backward_match_reference(durations, sigma, reduction): assert_grad_close(dur.grad, ref_dur.grad, atol=1e-5, rtol=1e-4) +@pytest.mark.kernels_ci def test_joint_logits_slices(): """Non-contiguous slices of a joint output (as in ParakeetForTDT) are handled without a copy.""" V = 128 @@ -258,6 +256,7 @@ def test_joint_logits_slices(): assert_grad_close(joint.grad, ref_joint.grad, atol=1e-5, rtol=1e-4) +@pytest.mark.kernels_ci @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_half_precision(dtype): """Half-precision logits are read directly and gradients come back in the input dtype.""" @@ -327,6 +326,7 @@ def test_long_targets(): assert_grad_close(tok.grad, ref_tok.grad, atol=GRAD_ATOL_LONG, rtol=0) +@pytest.mark.kernels_ci def test_empty_targets(): """Samples without labels: the only path is a sequence of blanks.""" token_logits, duration_logits, _, logit_lengths, _ = make_inputs(2, 10, 0, 16, DURATIONS) @@ -338,6 +338,7 @@ def test_empty_targets(): torch.testing.assert_close(losses.double().cpu(), expected, atol=1e-4, rtol=1e-5) +@pytest.mark.kernels_ci def test_infeasible_alignment(): """No valid path: the loss is infinite and the gradient is zero rather than NaN.""" durations = [0, 4] @@ -354,6 +355,7 @@ def test_infeasible_alignment(): assert (tok.grad[1] == 0).all() +@pytest.mark.kernels_ci def test_deterministic(): inputs = make_inputs(4, 30, 8, 64, DURATIONS) grads = [] @@ -364,6 +366,7 @@ def test_deterministic(): assert torch.equal(grads[0], grads[1]) +@pytest.mark.kernels_ci def test_invalid_reduction(): inputs = make_inputs(1, 5, 2, 4, [0, 1]) with pytest.raises(ValueError, match="Invalid reduction"): From 850114b3580b6f64fb2da479e09d3b898b5029a8 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 28 Sep 2026 17:37:24 +1000 Subject: [PATCH 09/17] tdt-loss: validate inputs, add launch bounds, share the test reference - Device-side asserts for out-of-range lengths and targets (like PyTorch indexing), and reject negative durations. - __launch_bounds__ on all kernels so the largest block always launches. - Cache the durations tensor per device instead of copying it every call. - Drop the unused log-likelihood output of the beta recursion. - Move the PyTorch reference to tests/reference.py, used by the benchmark too. Co-Authored-By: Claude Opus 5.5 --- tdt-loss/benchmarks/benchmark.py | 76 ++------------------- tdt-loss/tdt_loss/common.cuh | 22 ++++-- tdt-loss/tdt_loss/lattice.cu | 55 +++++++-------- tdt-loss/tdt_loss/logprobs.cu | 7 +- tdt-loss/tests/reference.py | 81 ++++++++++++++++++++++ tdt-loss/tests/test_tdt_loss.py | 89 +++---------------------- tdt-loss/torch-ext/tdt_loss/__init__.py | 27 +++++--- tdt-loss/torch-ext/torch_binding.cpp | 2 +- tdt-loss/torch-ext/torch_binding.h | 5 +- 9 files changed, 168 insertions(+), 196 deletions(-) create mode 100644 tdt-loss/tests/reference.py diff --git a/tdt-loss/benchmarks/benchmark.py b/tdt-loss/benchmarks/benchmark.py index 8d28cfbd1..edcd6f635 100644 --- a/tdt-loss/benchmarks/benchmark.py +++ b/tdt-loss/benchmarks/benchmark.py @@ -1,78 +1,14 @@ +import sys +from pathlib import Path + import torch from kernels.benchmark import Benchmark -DURATIONS = [0, 1, 2, 3, 4] - +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from tests.reference import tdt_loss_reference # noqa: E402 -def tdt_loss_reference( - token_logits, duration_logits, targets, logit_lengths, target_lengths, blank_id, durations, sigma=0.0 -): - """Per-sample TDT loss, as implemented in `transformers.loss.loss_tdt.tdt_loss`.""" - device = token_logits.device - batch_size, max_t, max_u, _ = token_logits.shape - - token_log_probs = torch.log_softmax(token_logits.float(), dim=-1) - sigma - duration_log_probs = torch.log_softmax(duration_logits.float(), dim=-1) - - log_alpha = torch.full((batch_size, max_t, max_u), float("-inf"), device=device) - log_alpha[:, 0, 0] = 0.0 - blank_log_probs = token_log_probs[:, :, :, blank_id] - if max_u > 1: - targets_expanded = targets.long().unsqueeze(1).expand(-1, max_t, -1) - label_log_probs = torch.gather( - token_log_probs[:, :, : max_u - 1, :], dim=3, index=targets_expanded.unsqueeze(-1) - ).squeeze(-1) - - neg_inf = torch.tensor(float("-inf"), device=device) - for n in range(1, max_t + max_u - 1): - u_indices = torch.arange(max(0, n - max_t + 1), min(n + 1, max_u), device=device) - t_indices = n - u_indices - candidates = [] - for i, dur in enumerate(durations): - t_prev = t_indices - dur - valid_t = t_prev >= 0 - if not valid_t.any(): - continue - t_src = t_prev.clamp(min=0) - if dur > 0: - contrib = ( - log_alpha[:, t_src, u_indices] - + blank_log_probs[:, t_src, u_indices] - + duration_log_probs[:, t_src, u_indices, i] - ) - candidates.append(torch.where(valid_t.unsqueeze(0), contrib, neg_inf)) - valid_both = valid_t & (u_indices > 0) - if valid_both.any(): - u_src = (u_indices - 1).clamp(min=0) - u_src_label = u_src.clamp(max=max_u - 2) if max_u > 1 else u_src - contrib = ( - log_alpha[:, t_src, u_src] - + label_log_probs[:, t_src, u_src_label] - + duration_log_probs[:, t_src, u_src, i] - ) - candidates.append(torch.where(valid_both.unsqueeze(0), contrib, neg_inf)) - if candidates: - log_alpha[:, t_indices, u_indices] = torch.logsumexp(torch.stack(candidates), dim=0) - - batch_idx = torch.arange(batch_size, device=device) - target_lengths = target_lengths.long() - log_probs = torch.full((batch_size,), float("-inf"), device=device) - for i, dur in enumerate(durations): - if dur == 0: - continue - t_final = logit_lengths.long() - dur - valid = t_final >= 0 - if not valid.any(): - continue - t_clamped = t_final.clamp(min=0) - terminal = ( - log_alpha[batch_idx, t_clamped, target_lengths] - + token_log_probs[batch_idx, t_clamped, target_lengths, blank_id] - + duration_log_probs[batch_idx, t_clamped, target_lengths, i] - ) - log_probs = torch.where(valid, torch.logsumexp(torch.stack([log_probs, terminal]), dim=0), log_probs) - return -log_probs +DURATIONS = [0, 1, 2, 3, 4] class TDTLossBenchmark(Benchmark): diff --git a/tdt-loss/tdt_loss/common.cuh b/tdt-loss/tdt_loss/common.cuh index 70d2d36ae..735a41c51 100644 --- a/tdt-loss/tdt_loss/common.cuh +++ b/tdt-loss/tdt_loss/common.cuh @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -10,6 +11,11 @@ namespace tdt_loss { constexpr int kWarpSize = 32; // Upper bound on the number of durations, only used to size shared memory. constexpr int kMaxDurations = 64; +// Block sizes of the per-node kernels and of the lattice recursions. The +// kernels are compiled with matching __launch_bounds__ so that register +// allocation always allows launching the largest block. +constexpr int kMaxRowThreads = 256; +constexpr int kMaxLatticeThreads = 1024; __device__ __forceinline__ float neg_inf() { return -INFINITY; } @@ -87,15 +93,21 @@ __device__ __forceinline__ void block_reduce_max_sum(float &m, float &s, s = shm_s[0]; } -// Clamp the per-sample lengths to the padded lattice so that malformed inputs -// cannot cause out-of-bounds accesses. Returns T (frames) and U (labels); the -// lattice for the sample spans t in [0, T) and u in [0, U]. +// Per-sample lengths: T (frames) and U (labels); the lattice for the sample +// spans t in [0, T) and u in [0, U]. Lengths outside of the padded lattice +// trigger a device-side assert, like out-of-range indexing does in PyTorch; +// they are also clamped so that no out-of-bounds access happens when asserts +// are compiled out. __device__ __forceinline__ void sample_lengths(const int *source_lengths, const int *target_lengths, int b, int max_T, int max_U, int &T, int &U) { - T = min(max(source_lengths[b], 0), max_T); - U = min(max(target_lengths[b], 0), max_U - 1); + T = source_lengths[b]; + U = target_lengths[b]; + CUDA_KERNEL_ASSERT(T >= 0 && T <= max_T && "source_lengths out of range"); + CUDA_KERNEL_ASSERT(U >= 0 && U < max_U && "target_lengths out of range"); + T = min(max(T, 0), max_T); + U = min(max(U, 0), max_U - 1); } } // namespace tdt_loss diff --git a/tdt-loss/tdt_loss/lattice.cu b/tdt-loss/tdt_loss/lattice.cu index be3eed1fa..af80f0001 100644 --- a/tdt-loss/tdt_loss/lattice.cu +++ b/tdt-loss/tdt_loss/lattice.cu @@ -32,7 +32,8 @@ __device__ __forceinline__ void load_durations(const int *durations, int D, __syncthreads(); } -__global__ void tdt_alpha_kernel(const float *__restrict__ blank_lp, +__global__ void __launch_bounds__(kMaxLatticeThreads) + tdt_alpha_kernel(const float *__restrict__ blank_lp, const float *__restrict__ label_lp, const float *__restrict__ dur_lp, const int *__restrict__ source_lengths, @@ -93,14 +94,14 @@ __global__ void tdt_alpha_kernel(const float *__restrict__ blank_lp, } } -__global__ void tdt_beta_kernel(const float *__restrict__ blank_lp, - const float *__restrict__ label_lp, - const float *__restrict__ dur_lp, - const int *__restrict__ source_lengths, - const int *__restrict__ target_lengths, - const int *__restrict__ durations, int max_T, - int max_U, int D, double *__restrict__ betas, - double *__restrict__ log_ll) { +__global__ void __launch_bounds__(kMaxLatticeThreads) + tdt_beta_kernel(const float *__restrict__ blank_lp, + const float *__restrict__ label_lp, + const float *__restrict__ dur_lp, + const int *__restrict__ source_lengths, + const int *__restrict__ target_lengths, + const int *__restrict__ durations, int max_T, int max_U, + int D, double *__restrict__ betas) { __shared__ int shm_durations[kMaxDurations]; const int b = blockIdx.x; @@ -143,15 +144,11 @@ __global__ void tdt_beta_kernel(const float *__restrict__ blank_lp, } __syncthreads(); } - - if (threadIdx.x == 0) { - log_ll[b] = T > 0 ? beta[0] : -INFINITY; - } } int lattice_block_size(int64_t max_U) { int threads = 32; - while (threads < 1024 && threads < max_U) threads *= 2; + while (threads < kMaxLatticeThreads && threads < max_U) threads *= 2; return threads; } @@ -161,10 +158,11 @@ void check_lattice_args(torch::Tensor const &blank_lp, torch::Tensor const &source_lengths, torch::Tensor const &target_lengths, torch::Tensor const &durations, - torch::Tensor const &out, torch::Tensor const &log_ll) { + torch::Tensor const &out, + torch::Tensor const *log_ll) { TORCH_CHECK(blank_lp.is_cuda(), "blank_lp must be a CUDA tensor"); TORCH_CHECK(blank_lp.dim() == 3, "blank_lp must have shape (batch, T, U+1)"); - for (auto const *x : {&blank_lp, &label_lp, &dur_lp, &out, &log_ll}) { + for (auto const *x : {&blank_lp, &label_lp, &dur_lp, &out}) { TORCH_CHECK(x->device() == blank_lp.device(), "all tensors must be on the same device"); TORCH_CHECK(x->is_contiguous(), "lattice tensors must be contiguous"); @@ -173,9 +171,8 @@ void check_lattice_args(torch::Tensor const &blank_lp, TORCH_CHECK(x->scalar_type() == torch::kFloat32, "log-prob tensors must be float32"); } - TORCH_CHECK(out.scalar_type() == torch::kFloat64 && - log_ll.scalar_type() == torch::kFloat64, - "alphas/betas and log_ll must be float64"); + TORCH_CHECK(out.scalar_type() == torch::kFloat64, + "alphas/betas must be float64"); for (auto const *x : {&source_lengths, &target_lengths, &durations}) { TORCH_CHECK(x->device() == blank_lp.device(), "all tensors must be on the same device"); @@ -189,9 +186,14 @@ void check_lattice_args(torch::Tensor const &blank_lp, "lattice tensors must have shape (batch, T, U+1)"); TORCH_CHECK(dur_lp.numel() == nodes * durations.numel(), "dur_lp must have shape (batch, T, U+1, num_durations)"); - TORCH_CHECK(source_lengths.numel() == B && target_lengths.numel() == B && - log_ll.numel() == B, - "lengths and log_ll must have shape (batch,)"); + TORCH_CHECK(source_lengths.numel() == B && target_lengths.numel() == B, + "lengths must have shape (batch,)"); + if (log_ll != nullptr) { + TORCH_CHECK(log_ll->device() == blank_lp.device() && + log_ll->scalar_type() == torch::kFloat64 && + log_ll->is_contiguous() && log_ll->numel() == B, + "log_ll must be a contiguous float64 tensor of shape (batch,)"); + } TORCH_CHECK(durations.numel() <= kMaxDurations, "at most ", kMaxDurations, " durations are supported"); } @@ -208,7 +210,7 @@ void tdt_loss_fwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, torch::Tensor &log_ll) { using namespace tdt_loss; check_lattice_args(blank_lp, label_lp, dur_lp, source_lengths, - target_lengths, durations, alphas, log_ll); + target_lengths, durations, alphas, &log_ll); const int64_t B = blank_lp.size(0); if (B == 0) return; @@ -227,11 +229,10 @@ 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 &log_ll) { + torch::Tensor const &durations, torch::Tensor &betas) { using namespace tdt_loss; check_lattice_args(blank_lp, label_lp, dur_lp, source_lengths, - target_lengths, durations, betas, log_ll); + target_lengths, durations, betas, nullptr); const int64_t B = blank_lp.size(0); if (B == 0) return; @@ -242,6 +243,6 @@ void tdt_loss_bwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, dur_lp.data_ptr(), source_lengths.data_ptr(), target_lengths.data_ptr(), durations.data_ptr(), blank_lp.size(1), blank_lp.size(2), durations.numel(), - betas.data_ptr(), log_ll.data_ptr()); + betas.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/tdt-loss/tdt_loss/logprobs.cu b/tdt-loss/tdt_loss/logprobs.cu index ec25f996a..bb54d7b44 100644 --- a/tdt-loss/tdt_loss/logprobs.cu +++ b/tdt-loss/tdt_loss/logprobs.cu @@ -18,7 +18,7 @@ namespace tdt_loss { namespace { template -__global__ void tdt_logprobs_fwd_kernel( +__global__ void __launch_bounds__(kMaxRowThreads) tdt_logprobs_fwd_kernel( const scalar_t *__restrict__ token_logits, int64_t tok_sb, int64_t tok_st, int64_t tok_su, const scalar_t *__restrict__ duration_logits, int64_t dur_sb, int64_t dur_st, int64_t dur_su, @@ -75,6 +75,7 @@ __global__ void tdt_logprobs_fwd_kernel( float label = neg_inf(); if (u < U) { const int target = targets[b * targets_stride + u]; + CUDA_KERNEL_ASSERT(target >= 0 && target < V && "target out of range"); if (target >= 0 && target < V) { label = static_cast(x[target]) - lse - sigma; } @@ -97,7 +98,7 @@ __global__ void tdt_logprobs_fwd_kernel( } template -__global__ void tdt_logits_grad_kernel( +__global__ void __launch_bounds__(kMaxRowThreads) tdt_logits_grad_kernel( const scalar_t *__restrict__ token_logits, int64_t tok_sb, int64_t tok_st, int64_t tok_su, const int *__restrict__ targets, int64_t targets_stride, const int *__restrict__ source_lengths, @@ -206,7 +207,7 @@ __global__ void tdt_logits_grad_kernel( int row_block_size(int64_t V) { int threads = 32; - while (threads < 256 && threads < V) threads *= 2; + while (threads < kMaxRowThreads && threads < V) threads *= 2; return threads; } diff --git a/tdt-loss/tests/reference.py b/tdt-loss/tests/reference.py new file mode 100644 index 000000000..62d6f470c --- /dev/null +++ b/tdt-loss/tests/reference.py @@ -0,0 +1,81 @@ +"""Pure PyTorch reference of the TDT loss, shared by the tests and the benchmark.""" + +import torch + + +def tdt_loss_reference( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + blank_id, + durations, + sigma=0.0, + dtype=torch.float32, +): + """Per-sample TDT loss, as implemented in `transformers.loss.loss_tdt.tdt_loss` (computed in `dtype`).""" + device = token_logits.device + batch_size, max_t, max_u, _ = token_logits.shape + + token_log_probs = torch.log_softmax(token_logits.to(dtype), dim=-1) - sigma + duration_log_probs = torch.log_softmax(duration_logits.to(dtype), dim=-1) + + log_alpha = torch.full((batch_size, max_t, max_u), float("-inf"), device=device, dtype=dtype) + log_alpha[:, 0, 0] = 0.0 + blank_log_probs = token_log_probs[:, :, :, blank_id] + if max_u > 1: + targets_expanded = targets.long().unsqueeze(1).expand(-1, max_t, -1) + label_log_probs = torch.gather( + token_log_probs[:, :, : max_u - 1, :], dim=3, index=targets_expanded.unsqueeze(-1) + ).squeeze(-1) + + neg_inf = torch.tensor(float("-inf"), device=device, dtype=dtype) + for n in range(1, max_t + max_u - 1): + u_indices = torch.arange(max(0, n - max_t + 1), min(n + 1, max_u), device=device) + t_indices = n - u_indices + candidates = [] + for i, dur in enumerate(durations): + t_prev = t_indices - dur + valid_t = t_prev >= 0 + if not valid_t.any(): + continue + t_src = t_prev.clamp(min=0) + if dur > 0: + contrib = ( + log_alpha[:, t_src, u_indices] + + blank_log_probs[:, t_src, u_indices] + + duration_log_probs[:, t_src, u_indices, i] + ) + candidates.append(torch.where(valid_t.unsqueeze(0), contrib, neg_inf)) + valid_both = valid_t & (u_indices > 0) + if valid_both.any(): + u_src = (u_indices - 1).clamp(min=0) + u_src_label = u_src.clamp(max=max_u - 2) if max_u > 1 else u_src + contrib = ( + log_alpha[:, t_src, u_src] + + label_log_probs[:, t_src, u_src_label] + + duration_log_probs[:, t_src, u_src, i] + ) + candidates.append(torch.where(valid_both.unsqueeze(0), contrib, neg_inf)) + if candidates: + log_alpha[:, t_indices, u_indices] = torch.logsumexp(torch.stack(candidates), dim=0) + + batch_idx = torch.arange(batch_size, device=device) + target_lengths = target_lengths.long() + log_probs = torch.full((batch_size,), float("-inf"), device=device, dtype=dtype) + for i, dur in enumerate(durations): + if dur == 0: + continue + t_final = logit_lengths.long() - dur + valid = t_final >= 0 + if not valid.any(): + continue + t_clamped = t_final.clamp(min=0) + terminal = ( + log_alpha[batch_idx, t_clamped, target_lengths] + + token_log_probs[batch_idx, t_clamped, target_lengths, blank_id] + + duration_log_probs[batch_idx, t_clamped, target_lengths, i] + ) + log_probs = torch.where(valid, torch.logsumexp(torch.stack([log_probs, terminal]), dim=0), log_probs) + return -log_probs diff --git a/tdt-loss/tests/test_tdt_loss.py b/tdt-loss/tests/test_tdt_loss.py index 8bf140267..c6fcc2dd0 100644 --- a/tdt-loss/tests/test_tdt_loss.py +++ b/tdt-loss/tests/test_tdt_loss.py @@ -1,7 +1,7 @@ """Tests for the TDT loss kernel. The kernel is checked against two pure PyTorch references: -- `tdt_loss_reference`: the vectorized implementation used in transformers, whose autograd +- `tdt_loss_reference` (tests/reference.py): the vectorized implementation used in transformers, whose autograd gradients are used to validate the kernel backward; - `tdt_loss_naive`: a loop-based implementation that is easy to verify by hand. """ @@ -11,6 +11,8 @@ import kernels +from tests.reference import tdt_loss_reference + if not torch.cuda.is_available(): pytest.skip("CUDA not available", allow_module_level=True) @@ -36,84 +38,6 @@ def reduce(losses, target_lengths, reduction): return losses -def tdt_loss_reference( - token_logits, - duration_logits, - targets, - logit_lengths, - target_lengths, - blank_id, - durations, - sigma=0.0, - dtype=torch.float32, -): - """Per-sample TDT loss, as implemented in `transformers.loss.loss_tdt.tdt_loss` (computed in `dtype`).""" - device = token_logits.device - batch_size, max_t, max_u, _ = token_logits.shape - - token_log_probs = torch.log_softmax(token_logits.to(dtype), dim=-1) - sigma - duration_log_probs = torch.log_softmax(duration_logits.to(dtype), dim=-1) - - log_alpha = torch.full((batch_size, max_t, max_u), float("-inf"), device=device, dtype=dtype) - log_alpha[:, 0, 0] = 0.0 - blank_log_probs = token_log_probs[:, :, :, blank_id] - if max_u > 1: - targets_expanded = targets.long().unsqueeze(1).expand(-1, max_t, -1) - label_log_probs = torch.gather( - token_log_probs[:, :, : max_u - 1, :], dim=3, index=targets_expanded.unsqueeze(-1) - ).squeeze(-1) - - neg_inf = torch.tensor(float("-inf"), device=device, dtype=dtype) - for n in range(1, max_t + max_u - 1): - u_indices = torch.arange(max(0, n - max_t + 1), min(n + 1, max_u), device=device) - t_indices = n - u_indices - candidates = [] - for i, dur in enumerate(durations): - t_prev = t_indices - dur - valid_t = t_prev >= 0 - if not valid_t.any(): - continue - t_src = t_prev.clamp(min=0) - if dur > 0: - contrib = ( - log_alpha[:, t_src, u_indices] - + blank_log_probs[:, t_src, u_indices] - + duration_log_probs[:, t_src, u_indices, i] - ) - candidates.append(torch.where(valid_t.unsqueeze(0), contrib, neg_inf)) - valid_both = valid_t & (u_indices > 0) - if valid_both.any(): - u_src = (u_indices - 1).clamp(min=0) - u_src_label = u_src.clamp(max=max_u - 2) if max_u > 1 else u_src - contrib = ( - log_alpha[:, t_src, u_src] - + label_log_probs[:, t_src, u_src_label] - + duration_log_probs[:, t_src, u_src, i] - ) - candidates.append(torch.where(valid_both.unsqueeze(0), contrib, neg_inf)) - if candidates: - log_alpha[:, t_indices, u_indices] = torch.logsumexp(torch.stack(candidates), dim=0) - - batch_idx = torch.arange(batch_size, device=device) - target_lengths = target_lengths.long() - log_probs = torch.full((batch_size,), float("-inf"), device=device, dtype=dtype) - for i, dur in enumerate(durations): - if dur == 0: - continue - t_final = logit_lengths.long() - dur - valid = t_final >= 0 - if not valid.any(): - continue - t_clamped = t_final.clamp(min=0) - terminal = ( - log_alpha[batch_idx, t_clamped, target_lengths] - + token_log_probs[batch_idx, t_clamped, target_lengths, blank_id] - + duration_log_probs[batch_idx, t_clamped, target_lengths, i] - ) - log_probs = torch.where(valid, torch.logsumexp(torch.stack([log_probs, terminal]), dim=0), log_probs) - return -log_probs - - def tdt_loss_naive(token_logits, duration_logits, targets, logit_lengths, target_lengths, blank_id, durations, sigma=0.0): """Per-sample TDT loss with explicit loops over the lattice, in float64 on the CPU.""" token_lp = (torch.log_softmax(token_logits.double().cpu(), dim=-1) - sigma).tolist() @@ -371,3 +295,10 @@ def test_invalid_reduction(): inputs = make_inputs(1, 5, 2, 4, [0, 1]) with pytest.raises(ValueError, match="Invalid reduction"): run_kernel(*inputs, [0, 1], 0.0, "invalid") + + +@pytest.mark.kernels_ci +def test_negative_durations(): + inputs = make_inputs(1, 5, 2, 4, [0, 1]) + with pytest.raises(ValueError, match="non-negative"): + run_kernel(*inputs, [-1, 1], 0.0, "mean") diff --git a/tdt-loss/torch-ext/tdt_loss/__init__.py b/tdt-loss/torch-ext/tdt_loss/__init__.py index 841bab359..5733f6c56 100644 --- a/tdt-loss/torch-ext/tdt_loss/__init__.py +++ b/tdt-loss/torch-ext/tdt_loss/__init__.py @@ -1,6 +1,6 @@ """TDT (Token-and-Duration Transducer) loss CUDA kernel.""" -from typing import Sequence, Union +from typing import Dict, Sequence, Tuple, Union import torch @@ -11,6 +11,21 @@ _REDUCTIONS = ("mean_volume", "mean_batch", "mean", "sum", "none") +_DURATIONS_CACHE: Dict[Tuple[Tuple[int, ...], torch.device], torch.Tensor] = {} + + +def _durations_tensor(durations: Union[Sequence[int], torch.Tensor], device: torch.device) -> torch.Tensor: + """Validated int32 tensor of durations on `device`, cached since the durations are fixed for a model.""" + if isinstance(durations, torch.Tensor): + durations = durations.tolist() + key = (tuple(int(d) for d in durations), device) + if key not in _DURATIONS_CACHE: + if any(d < 0 for d in key[0]): + raise ValueError(f"Durations must be non-negative, got {list(key[0])}.") + _DURATIONS_CACHE[key] = torch.tensor(key[0], device=device, dtype=torch.int32) + return _DURATIONS_CACHE[key] + + def _last_dim_contiguous(x: torch.Tensor) -> torch.Tensor: # The kernels accept arbitrary strides for the (batch, T, U) dims, so slices of a # joint `(..., vocab_size + num_durations)` output can be passed without a copy. @@ -100,8 +115,7 @@ def backward(ctx, grad_loss: torch.Tensor): return (None,) * 8 betas = torch.empty_like(alphas) - log_ll_bwd = torch.empty_like(log_ll) - ops.tdt_loss_bwd(blank_lp, label_lp, dur_lp, source_lengths, target_lengths, durations, betas, log_ll_bwd) + ops.tdt_loss_bwd(blank_lp, label_lp, dur_lp, source_lengths, target_lengths, durations, betas) grad_token_logits = torch.empty(token_logits.shape, device=token_logits.device, dtype=token_logits.dtype) grad_duration_logits = torch.empty(ctx.duration_shape, device=token_logits.device, dtype=ctx.duration_dtype) @@ -146,7 +160,7 @@ def tdt_loss( targets: Target labels of shape `(batch, U)`. source_lengths: Encoder output lengths of shape `(batch,)`. target_lengths: Target lengths of shape `(batch,)`. - durations: Duration values, e.g. `[0, 1, 2, 3, 4]`. + durations: Non-negative 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: One of `"mean_volume"`, `"mean_batch"`, `"mean"`, `"sum"` or `"none"`, mirroring NeMo's @@ -162,10 +176,7 @@ def tdt_loss( ) device = token_logits.device - if isinstance(durations, torch.Tensor): - durations = durations.to(device=device, dtype=torch.int32).contiguous() - else: - durations = torch.tensor(list(durations), device=device, dtype=torch.int32) + durations = _durations_tensor(durations, device) targets = targets.to(device=device, dtype=torch.int32).contiguous() source_lengths = source_lengths.to(device=device, dtype=torch.int32).contiguous() target_lengths = target_lengths.to(device=device, dtype=torch.int32).contiguous() diff --git a/tdt-loss/torch-ext/torch_binding.cpp b/tdt-loss/torch-ext/torch_binding.cpp index f19d41276..b3d91f420 100644 --- a/tdt-loss/torch-ext/torch_binding.cpp +++ b/tdt-loss/torch-ext/torch_binding.cpp @@ -20,7 +20,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { 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! log_ll) -> ()"); + "Tensor! betas) -> ()"); ops.impl("tdt_loss_bwd", torch::kCUDA, &tdt_loss_bwd); ops.def( diff --git a/tdt-loss/torch-ext/torch_binding.h b/tdt-loss/torch-ext/torch_binding.h index d9c7d32d1..c1337b4ec 100644 --- a/tdt-loss/torch-ext/torch_binding.h +++ b/tdt-loss/torch-ext/torch_binding.h @@ -22,13 +22,12 @@ void tdt_loss_fwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, torch::Tensor const &durations, torch::Tensor &alphas, torch::Tensor &log_ll); -// Backward recursion over the lattice (betas) and per-sample log-likelihood. +// Backward recursion over the lattice (betas). 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 &log_ll); + torch::Tensor const &durations, torch::Tensor &betas); // Gradient of the per-sample losses (scaled by grad_loss) w.r.t. the token // and duration logits. From f711276734dc1768eb4043e84f79d61b9e4925ed Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 28 Sep 2026 20:10:33 +1000 Subject: [PATCH 10/17] tdt-loss: add A100/H200 benchmarks and fine-tuning results to the README Co-Authored-By: Claude Opus 5.5 --- tdt-loss/README.md | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/tdt-loss/README.md b/tdt-loss/README.md index 02fb59782..9e5284d1f 100644 --- a/tdt-loss/README.md +++ b/tdt-loss/README.md @@ -59,14 +59,22 @@ Samples without a valid alignment get an infinite loss and a zero gradient. ## Benchmarks -Forward + backward on an NVIDIA L4, with the logits given as slices of a joint output, compared with the vectorized +Forward + backward of the loss, with the logits given as slices of a joint output, compared with the vectorized PyTorch implementation from `transformers.loss.loss_tdt`. Run `benchmarks/benchmark.py` with `kernels benchmark`. -| batch, frames, labels, vocab | dtype | PyTorch | kernel | speedup | -|---|---|---:|---:|---:| -| 8, 100, 20, 1025 | float32 | 536 ms | 4.3 ms | 125x | -| 4, 200, 40, 8193 | float32 | 1252 ms | 53 ms | 24x | -| 4, 200, 40, 8193 | bfloat16 | 1242 ms | 39 ms | 32x | +| batch, frames, labels, vocab | dtype | GPU | PyTorch | kernel | speedup | +|---|---|---|---:|---:|---:| +| 8, 100, 20, 1025 | float32 | L4 | 536 ms | 4.3 ms | 125x | +| 4, 200, 40, 8193 | float32 | L4 | 1252 ms | 53 ms | 24x | +| 4, 200, 40, 8193 | bfloat16 | L4 | 1242 ms | 39 ms | 32x | +| 8, 100, 20, 1025 | float32 | A100 | 624 ms | 1.5 ms | 410x | +| 4, 200, 40, 8193 | float32 | A100 | 1259 ms | 9.5 ms | 133x | +| 4, 200, 40, 8193 | bfloat16 | A100 | 1267 ms | 8.6 ms | 148x | +| 8, 100, 20, 1025 | float32 | H200 | 488 ms | 1.0 ms | 480x | +| 4, 200, 40, 8193 | float32 | H200 | 993 ms | 5.0 ms | 199x | +| 4, 200, 40, 8193 | bfloat16 | H200 | 896 ms | 4.9 ms | 183x | -The kernels run close to the memory bandwidth of the GPU; most of the remaining time is spent by PyTorch -materializing the gradient of the two logits slices. +Fine-tuning [nvidia/parakeet-tdt-0.6b-v3](https://huggingface.co/nvidia/parakeet-tdt-0.6b-v3) on LibriSpeech +(batch size 8, bfloat16 autocast, 300 steps, same data order and seed) on an H200 gives the same loss curve as the +PyTorch loss (identical first-step loss, median per-step difference of 0.0015), with a full training step going from +1305 ms to 144 ms (9x) and peak memory from 33.2 GiB to 21.5 GiB. From cfeecb25b220f1dc15cfbb4b589ddb92a8c90e07 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Tue, 29 Sep 2026 00:31:52 +1000 Subject: [PATCH 11/17] tdt-loss: register in the maintainer registry and freshness source mapping Co-Authored-By: Claude Opus 5.5 --- .github/kernel-maintainers.json | 1 + scripts/check_kernel_freshness.py | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/kernel-maintainers.json b/.github/kernel-maintainers.json index 16eacacd4..923a929b6 100644 --- a/.github/kernel-maintainers.json +++ b/.github/kernel-maintainers.json @@ -51,6 +51,7 @@ "scattermoe": ["U08Q42K8MGW"], "sgl-flash-attn3": [], "sonic-moe": ["U08Q42K8MGW", "U054KJK89C2"], + "tdt-loss": [], "tinygrad-rms": [], "topk": ["U0586KAH2E9"], "trimul-gpumode": [], diff --git a/scripts/check_kernel_freshness.py b/scripts/check_kernel_freshness.py index 47b59c5cb..20302a861 100644 --- a/scripts/check_kernel_freshness.py +++ b/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", From 3b9907c58a98e3874b16b454a72fe8fb2a8b67c7 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Tue, 29 Sep 2026 00:44:36 +1000 Subject: [PATCH 12/17] tdt-loss: add a TDTLoss layer for the kernels exchange pattern Move the implementation from __init__.py to loss.py, add layers.TDTLoss (a stateless layer with the same signature as transformers' tdt_loss) and use the transformers argument order in the functional API. --- tdt-loss/README.md | 5 +- tdt-loss/benchmarks/benchmark.py | 2 +- tdt-loss/tests/test_tdt_loss.py | 36 ++++- tdt-loss/torch-ext/tdt_loss/__init__.py | 203 +----------------------- tdt-loss/torch-ext/tdt_loss/layers.py | 40 +++++ tdt-loss/torch-ext/tdt_loss/loss.py | 201 +++++++++++++++++++++++ 6 files changed, 284 insertions(+), 203 deletions(-) create mode 100644 tdt-loss/torch-ext/tdt_loss/layers.py create mode 100644 tdt-loss/torch-ext/tdt_loss/loss.py diff --git a/tdt-loss/README.md b/tdt-loss/README.md index 9e5284d1f..3bd36fddc 100644 --- a/tdt-loss/README.md +++ b/tdt-loss/README.md @@ -47,8 +47,8 @@ loss = tdt.tdt_loss( targets, logit_lengths, target_lengths, - durations, blank_id, + durations, sigma=0.0, reduction="mean", # or "mean_volume", "mean_batch", "sum", "none" ) @@ -57,6 +57,9 @@ loss.backward() Samples without a valid alignment get an infinite loss and a zero gradient. +The `layers.TDTLoss` layer has the same signature as `transformers.loss.loss_tdt.tdt_loss`, so 🤗 Transformers +swaps it in for Parakeet TDT models loaded with `use_kernels=True`. + ## Benchmarks Forward + backward of the loss, with the logits given as slices of a joint output, compared with the vectorized diff --git a/tdt-loss/benchmarks/benchmark.py b/tdt-loss/benchmarks/benchmark.py index edcd6f635..42fffaa34 100644 --- a/tdt-loss/benchmarks/benchmark.py +++ b/tdt-loss/benchmarks/benchmark.py @@ -31,8 +31,8 @@ def _run(self): self.targets, self.logit_lengths, self.target_lengths, - DURATIONS, 0, + DURATIONS, reduction="none", ) losses.sum().backward() diff --git a/tdt-loss/tests/test_tdt_loss.py b/tdt-loss/tests/test_tdt_loss.py index c6fcc2dd0..a0af3d3a6 100644 --- a/tdt-loss/tests/test_tdt_loss.py +++ b/tdt-loss/tests/test_tdt_loss.py @@ -106,8 +106,8 @@ def run_kernel(token_logits, duration_logits, targets, logit_lengths, target_len targets, logit_lengths, target_lengths, - durations, 0, + durations, sigma=sigma, reduction=reduction, ) @@ -302,3 +302,37 @@ def test_negative_durations(): inputs = make_inputs(1, 5, 2, 4, [0, 1]) with pytest.raises(ValueError, match="non-negative"): run_kernel(*inputs, [-1, 1], 0.0, "mean") + + +@pytest.mark.kernels_ci +def test_layer(): + """The `TDTLoss` layer matches the functional API and mirrors the `transformers` `tdt_loss` signature.""" + import inspect + + params = list(inspect.signature(tdt_loss_kernel.layers.TDTLoss.forward).parameters) + assert params == [ + "self", + "token_logits", + "duration_logits", + "targets", + "logit_lengths", + "target_lengths", + "blank_token_id", + "durations", + "sigma", + "reduction", + ] + + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs(2, 12, 4, 16, DURATIONS) + kwargs = dict( + token_logits=token_logits, + duration_logits=duration_logits, + targets=targets, + logit_lengths=logit_lengths, + target_lengths=target_lengths, + blank_token_id=0, + durations=DURATIONS, + sigma=0.05, + reduction="none", + ) + torch.testing.assert_close(tdt_loss_kernel.layers.TDTLoss()(**kwargs), tdt_loss_kernel.tdt_loss(**kwargs)) diff --git a/tdt-loss/torch-ext/tdt_loss/__init__.py b/tdt-loss/torch-ext/tdt_loss/__init__.py index 5733f6c56..e5f7d3cb9 100644 --- a/tdt-loss/torch-ext/tdt_loss/__init__.py +++ b/tdt-loss/torch-ext/tdt_loss/__init__.py @@ -1,203 +1,6 @@ """TDT (Token-and-Duration Transducer) loss CUDA kernel.""" -from typing import Dict, Sequence, Tuple, Union +from . import layers +from .loss import tdt_loss -import torch - -from ._ops import ops - -__all__ = ["tdt_loss"] - -_REDUCTIONS = ("mean_volume", "mean_batch", "mean", "sum", "none") - - -_DURATIONS_CACHE: Dict[Tuple[Tuple[int, ...], torch.device], torch.Tensor] = {} - - -def _durations_tensor(durations: Union[Sequence[int], torch.Tensor], device: torch.device) -> torch.Tensor: - """Validated int32 tensor of durations on `device`, cached since the durations are fixed for a model.""" - if isinstance(durations, torch.Tensor): - durations = durations.tolist() - key = (tuple(int(d) for d in durations), device) - if key not in _DURATIONS_CACHE: - if any(d < 0 for d in key[0]): - raise ValueError(f"Durations must be non-negative, got {list(key[0])}.") - _DURATIONS_CACHE[key] = torch.tensor(key[0], device=device, dtype=torch.int32) - return _DURATIONS_CACHE[key] - - -def _last_dim_contiguous(x: torch.Tensor) -> torch.Tensor: - # The kernels accept arbitrary strides for the (batch, T, U) dims, so slices of a - # joint `(..., vocab_size + num_durations)` output can be passed without a copy. - return x if x.stride(-1) == 1 else x.contiguous() - - -class TDTLoss(torch.autograd.Function): - """Per-sample TDT loss (negative log-likelihood) with a CUDA forward and backward.""" - - @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, - ) -> torch.Tensor: - token_logits = _last_dim_contiguous(token_logits) - duration_logits = _last_dim_contiguous(duration_logits) - - B, max_T, max_U, _ = token_logits.shape - D = duration_logits.shape[-1] - f32 = dict(device=token_logits.device, dtype=torch.float32) - - blank_lp = torch.empty(B, max_T, max_U, **f32) - label_lp = torch.empty(B, max_T, max_U, **f32) - dur_lp = torch.empty(B, max_T, max_U, D, **f32) - token_lse = torch.empty(B, max_T, max_U, **f32) - ops.tdt_logprobs_fwd( - token_logits, - duration_logits, - targets, - source_lengths, - target_lengths, - blank_id, - sigma, - blank_lp, - label_lp, - dur_lp, - token_lse, - ) - - # The lattice recursions run in float64, see lattice.cu. - f64 = dict(device=token_logits.device, dtype=torch.float64) - alphas = torch.empty(B, max_T, max_U, **f64) - log_ll = torch.empty(B, **f64) - ops.tdt_loss_fwd(blank_lp, label_lp, dur_lp, source_lengths, target_lengths, durations, alphas, log_ll) - - ctx.save_for_backward( - token_logits, - targets, - source_lengths, - target_lengths, - durations, - blank_lp, - label_lp, - dur_lp, - token_lse, - alphas, - log_ll, - ) - ctx.blank_id = blank_id - ctx.duration_shape = duration_logits.shape - ctx.duration_dtype = duration_logits.dtype - return -log_ll.float() - - @staticmethod - def backward(ctx, grad_loss: torch.Tensor): - ( - token_logits, - targets, - source_lengths, - target_lengths, - durations, - blank_lp, - label_lp, - dur_lp, - token_lse, - alphas, - log_ll, - ) = ctx.saved_tensors - if not (ctx.needs_input_grad[0] or ctx.needs_input_grad[1]): - return (None,) * 8 - - betas = torch.empty_like(alphas) - ops.tdt_loss_bwd(blank_lp, label_lp, dur_lp, source_lengths, target_lengths, durations, betas) - - grad_token_logits = torch.empty(token_logits.shape, device=token_logits.device, dtype=token_logits.dtype) - grad_duration_logits = torch.empty(ctx.duration_shape, device=token_logits.device, dtype=ctx.duration_dtype) - ops.tdt_logits_grad( - token_logits, - targets, - source_lengths, - target_lengths, - durations, - blank_lp, - label_lp, - dur_lp, - token_lse, - alphas, - betas, - log_ll, - grad_loss.float().contiguous(), - ctx.blank_id, - 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[Sequence[int], torch.Tensor], - blank_id: int, - sigma: float = 0.0, - reduction: str = "mean", -) -> torch.Tensor: - """Compute the TDT (Token-and-Duration Transducer) loss (https://arxiv.org/abs/2304.06795). - - Args: - token_logits: Token logits of shape `(batch, T, U+1, vocab_size+1)`, in float32, float16 or bfloat16. - Only the last dimension needs to be contiguous. - duration_logits: Duration logits of shape `(batch, T, U+1, num_durations)`, same dtype as `token_logits`. - targets: Target labels of shape `(batch, U)`. - source_lengths: Encoder output lengths of shape `(batch,)`. - target_lengths: Target lengths of shape `(batch,)`. - durations: Non-negative 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: One of `"mean_volume"`, `"mean_batch"`, `"mean"`, `"sum"` or `"none"`, mirroring NeMo's - `RNNTLoss`. `"mean"` divides each loss by its target length before averaging over the batch. - - Returns: - Scalar loss tensor (or per-example losses of shape `(batch,)` if `reduction="none"`). Samples without a - valid alignment get an infinite loss and a zero gradient. - """ - if reduction not in _REDUCTIONS: - raise ValueError( - f'Invalid reduction mode "{reduction}". Expected one of {", ".join(repr(r) for r in _REDUCTIONS)}.' - ) - - device = token_logits.device - durations = _durations_tensor(durations, device) - targets = targets.to(device=device, dtype=torch.int32).contiguous() - source_lengths = source_lengths.to(device=device, dtype=torch.int32).contiguous() - target_lengths = target_lengths.to(device=device, dtype=torch.int32).contiguous() - - losses = TDTLoss.apply( - token_logits, - duration_logits, - targets, - source_lengths, - target_lengths, - durations, - int(blank_id), - float(sigma), - ) - - if reduction == "mean_volume": - return losses.sum() / target_lengths.sum().float() - if reduction == "mean_batch": - return losses.mean() - if reduction == "mean": - return (losses / target_lengths.float()).mean() - if reduction == "sum": - return losses.sum() - return losses +__all__ = ["layers", "tdt_loss"] diff --git a/tdt-loss/torch-ext/tdt_loss/layers.py b/tdt-loss/torch-ext/tdt_loss/layers.py new file mode 100644 index 000000000..fcede711c --- /dev/null +++ b/tdt-loss/torch-ext/tdt_loss/layers.py @@ -0,0 +1,40 @@ +from typing import Sequence, Union + +import torch +import torch.nn as nn + +from .loss import tdt_loss + + +class TDTLoss(nn.Module): + """ + Stateless layer computing the TDT loss with the CUDA kernels. Its forward has the same signature as + `transformers.loss.loss_tdt.tdt_loss`, so it can replace it through `kernels.kernelize`. + """ + + has_backward = True + can_torch_compile = False + + def forward( + self, + token_logits: torch.Tensor, + duration_logits: torch.Tensor, + targets: torch.Tensor, + logit_lengths: torch.Tensor, + target_lengths: torch.Tensor, + blank_token_id: int, + durations: Union[Sequence[int], torch.Tensor], + sigma: float = 0.0, + reduction: str = "mean", + ) -> torch.Tensor: + return tdt_loss( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + blank_token_id, + durations, + sigma=sigma, + reduction=reduction, + ) diff --git a/tdt-loss/torch-ext/tdt_loss/loss.py b/tdt-loss/torch-ext/tdt_loss/loss.py new file mode 100644 index 000000000..01f646501 --- /dev/null +++ b/tdt-loss/torch-ext/tdt_loss/loss.py @@ -0,0 +1,201 @@ +"""TDT (Token-and-Duration Transducer) loss: autograd function and functional API.""" + +from typing import Dict, Sequence, Tuple, Union + +import torch + +from ._ops import ops + +_REDUCTIONS = ("mean_volume", "mean_batch", "mean", "sum", "none") + + +_DURATIONS_CACHE: Dict[Tuple[Tuple[int, ...], torch.device], torch.Tensor] = {} + + +def _durations_tensor(durations: Union[Sequence[int], torch.Tensor], device: torch.device) -> torch.Tensor: + """Validated int32 tensor of durations on `device`, cached since the durations are fixed for a model.""" + if isinstance(durations, torch.Tensor): + durations = durations.tolist() + key = (tuple(int(d) for d in durations), device) + if key not in _DURATIONS_CACHE: + if any(d < 0 for d in key[0]): + raise ValueError(f"Durations must be non-negative, got {list(key[0])}.") + _DURATIONS_CACHE[key] = torch.tensor(key[0], device=device, dtype=torch.int32) + return _DURATIONS_CACHE[key] + + +def _last_dim_contiguous(x: torch.Tensor) -> torch.Tensor: + # The kernels accept arbitrary strides for the (batch, T, U) dims, so slices of a + # joint `(..., vocab_size + num_durations)` output can be passed without a copy. + return x if x.stride(-1) == 1 else x.contiguous() + + +class _TDTLossFunction(torch.autograd.Function): + """Per-sample TDT loss (negative log-likelihood) with a CUDA forward and backward.""" + + @staticmethod + def forward( + ctx, + token_logits: torch.Tensor, + duration_logits: torch.Tensor, + targets: torch.Tensor, + logit_lengths: torch.Tensor, + target_lengths: torch.Tensor, + durations: torch.Tensor, + blank_id: int, + sigma: float, + ) -> torch.Tensor: + token_logits = _last_dim_contiguous(token_logits) + duration_logits = _last_dim_contiguous(duration_logits) + + B, max_T, max_U, _ = token_logits.shape + D = duration_logits.shape[-1] + f32 = dict(device=token_logits.device, dtype=torch.float32) + + blank_lp = torch.empty(B, max_T, max_U, **f32) + label_lp = torch.empty(B, max_T, max_U, **f32) + dur_lp = torch.empty(B, max_T, max_U, D, **f32) + token_lse = torch.empty(B, max_T, max_U, **f32) + ops.tdt_logprobs_fwd( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + blank_id, + sigma, + blank_lp, + label_lp, + dur_lp, + token_lse, + ) + + # The lattice recursions run in float64, see lattice.cu. + f64 = dict(device=token_logits.device, dtype=torch.float64) + alphas = torch.empty(B, max_T, max_U, **f64) + log_ll = torch.empty(B, **f64) + ops.tdt_loss_fwd(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations, alphas, log_ll) + + ctx.save_for_backward( + token_logits, + targets, + logit_lengths, + target_lengths, + durations, + blank_lp, + label_lp, + dur_lp, + token_lse, + alphas, + log_ll, + ) + ctx.blank_id = blank_id + ctx.duration_shape = duration_logits.shape + ctx.duration_dtype = duration_logits.dtype + return -log_ll.float() + + @staticmethod + def backward(ctx, grad_loss: torch.Tensor): + ( + token_logits, + targets, + logit_lengths, + target_lengths, + durations, + blank_lp, + label_lp, + dur_lp, + token_lse, + alphas, + log_ll, + ) = ctx.saved_tensors + if not (ctx.needs_input_grad[0] or ctx.needs_input_grad[1]): + return (None,) * 8 + + betas = torch.empty_like(alphas) + ops.tdt_loss_bwd(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations, betas) + + grad_token_logits = torch.empty(token_logits.shape, device=token_logits.device, dtype=token_logits.dtype) + grad_duration_logits = torch.empty(ctx.duration_shape, device=token_logits.device, dtype=ctx.duration_dtype) + ops.tdt_logits_grad( + token_logits, + targets, + logit_lengths, + target_lengths, + durations, + blank_lp, + label_lp, + dur_lp, + token_lse, + alphas, + betas, + log_ll, + grad_loss.float().contiguous(), + ctx.blank_id, + 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, + logit_lengths: torch.Tensor, + target_lengths: torch.Tensor, + blank_token_id: int, + durations: Union[Sequence[int], torch.Tensor], + sigma: float = 0.0, + reduction: str = "mean", +) -> torch.Tensor: + """Compute the TDT (Token-and-Duration Transducer) loss (https://arxiv.org/abs/2304.06795). + + Args: + token_logits: Token logits of shape `(batch, T, U+1, vocab_size+1)`, in float32, float16 or bfloat16. + Only the last dimension needs to be contiguous. + duration_logits: Duration logits of shape `(batch, T, U+1, num_durations)`, same dtype as `token_logits`. + targets: Target labels of shape `(batch, U)`. + logit_lengths: Encoder output lengths of shape `(batch,)`. + target_lengths: Target lengths of shape `(batch,)`. + blank_token_id: Blank token id. + durations: Non-negative duration values, e.g. `[0, 1, 2, 3, 4]`. + sigma: Logit undernormalization constant (see TDT paper). Defaults to `0.0`. + reduction: One of `"mean_volume"`, `"mean_batch"`, `"mean"`, `"sum"` or `"none"`, mirroring NeMo's + `RNNTLoss`. `"mean"` divides each loss by its target length before averaging over the batch. + + Returns: + Scalar loss tensor (or per-example losses of shape `(batch,)` if `reduction="none"`). Samples without a + valid alignment get an infinite loss and a zero gradient. + """ + if reduction not in _REDUCTIONS: + raise ValueError( + f'Invalid reduction mode "{reduction}". Expected one of {", ".join(repr(r) for r in _REDUCTIONS)}.' + ) + + device = token_logits.device + durations = _durations_tensor(durations, device) + targets = targets.to(device=device, dtype=torch.int32).contiguous() + logit_lengths = logit_lengths.to(device=device, dtype=torch.int32).contiguous() + target_lengths = target_lengths.to(device=device, dtype=torch.int32).contiguous() + + losses = _TDTLossFunction.apply( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + durations, + int(blank_token_id), + float(sigma), + ) + + if reduction == "mean_volume": + return losses.sum() / target_lengths.sum().float() + if reduction == "mean_batch": + return losses.mean() + if reduction == "mean": + return (losses / target_lengths.float()).mean() + if reduction == "sum": + return losses.sum() + return losses From 02647aff4a7c7a6640823f9b298465844c86bb73 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Fri, 2 Oct 2026 11:49:51 +0200 Subject: [PATCH 13/17] tdt-loss: support torch.compile Register fake implementations for the ops (they only write to preallocated outputs), build the durations tensor without the eager cache while compiling, and mark the TDTLoss layer as compile-friendly. Test with torch.compile(fullgraph=True) and torch.library.opcheck. --- tdt-loss/README.md | 5 ++- tdt-loss/tdt_loss/common.cuh | 6 +-- tdt-loss/tdt_loss/lattice.cu | 26 ++++++------ tdt-loss/tdt_loss/logprobs.cu | 22 +++++----- tdt-loss/tests/test_tdt_loss.py | 57 +++++++++++++++++++++++++ tdt-loss/torch-ext/tdt_loss/layers.py | 2 +- tdt-loss/torch-ext/tdt_loss/loss.py | 61 +++++++++++++++++++++++++-- tdt-loss/torch-ext/torch_binding.cpp | 8 ++-- tdt-loss/torch-ext/torch_binding.h | 8 ++-- 9 files changed, 153 insertions(+), 42 deletions(-) diff --git a/tdt-loss/README.md b/tdt-loss/README.md index 3bd36fddc..e3cbb6314 100644 --- a/tdt-loss/README.md +++ b/tdt-loss/README.md @@ -57,8 +57,9 @@ loss.backward() Samples without a valid alignment get an infinite loss and a zero gradient. -The `layers.TDTLoss` layer has the same signature as `transformers.loss.loss_tdt.tdt_loss`, so 🤗 Transformers -swaps it in for Parakeet TDT models loaded with `use_kernels=True`. +The loss supports `torch.compile` (including `fullgraph=True`). The `layers.TDTLoss` layer has the same signature +as `transformers.loss.loss_tdt.tdt_loss`, so 🤗 Transformers swaps it in for Parakeet TDT models loaded with +`use_kernels=True`. ## Benchmarks diff --git a/tdt-loss/tdt_loss/common.cuh b/tdt-loss/tdt_loss/common.cuh index 735a41c51..de02e7ca1 100644 --- a/tdt-loss/tdt_loss/common.cuh +++ b/tdt-loss/tdt_loss/common.cuh @@ -98,13 +98,13 @@ __device__ __forceinline__ void block_reduce_max_sum(float &m, float &s, // trigger a device-side assert, like out-of-range indexing does in PyTorch; // they are also clamped so that no out-of-bounds access happens when asserts // are compiled out. -__device__ __forceinline__ void sample_lengths(const int *source_lengths, +__device__ __forceinline__ void sample_lengths(const int *logit_lengths, const int *target_lengths, int b, int max_T, int max_U, int &T, int &U) { - T = source_lengths[b]; + T = logit_lengths[b]; U = target_lengths[b]; - CUDA_KERNEL_ASSERT(T >= 0 && T <= max_T && "source_lengths out of range"); + CUDA_KERNEL_ASSERT(T >= 0 && T <= max_T && "logit_lengths out of range"); CUDA_KERNEL_ASSERT(U >= 0 && U < max_U && "target_lengths out of range"); T = min(max(T, 0), max_T); U = min(max(U, 0), max_U - 1); diff --git a/tdt-loss/tdt_loss/lattice.cu b/tdt-loss/tdt_loss/lattice.cu index af80f0001..27e1a94e0 100644 --- a/tdt-loss/tdt_loss/lattice.cu +++ b/tdt-loss/tdt_loss/lattice.cu @@ -36,7 +36,7 @@ __global__ void __launch_bounds__(kMaxLatticeThreads) tdt_alpha_kernel(const float *__restrict__ blank_lp, const float *__restrict__ label_lp, const float *__restrict__ dur_lp, - const int *__restrict__ source_lengths, + const int *__restrict__ logit_lengths, const int *__restrict__ target_lengths, const int *__restrict__ durations, int max_T, int max_U, int D, double *__restrict__ alphas, @@ -45,7 +45,7 @@ __global__ void __launch_bounds__(kMaxLatticeThreads) const int b = blockIdx.x; int T, U; - sample_lengths(source_lengths, target_lengths, b, max_T, max_U, T, U); + sample_lengths(logit_lengths, target_lengths, b, max_T, max_U, T, U); load_durations(durations, D, shm_durations); const int64_t sample = static_cast(b) * max_T * max_U; @@ -98,7 +98,7 @@ __global__ void __launch_bounds__(kMaxLatticeThreads) tdt_beta_kernel(const float *__restrict__ blank_lp, const float *__restrict__ label_lp, const float *__restrict__ dur_lp, - const int *__restrict__ source_lengths, + const int *__restrict__ logit_lengths, const int *__restrict__ target_lengths, const int *__restrict__ durations, int max_T, int max_U, int D, double *__restrict__ betas) { @@ -106,7 +106,7 @@ __global__ void __launch_bounds__(kMaxLatticeThreads) const int b = blockIdx.x; int T, U; - sample_lengths(source_lengths, target_lengths, b, max_T, max_U, T, U); + sample_lengths(logit_lengths, target_lengths, b, max_T, max_U, T, U); load_durations(durations, D, shm_durations); const int64_t sample = static_cast(b) * max_T * max_U; @@ -155,7 +155,7 @@ int lattice_block_size(int64_t max_U) { void check_lattice_args(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, torch::Tensor const &dur_lp, - torch::Tensor const &source_lengths, + torch::Tensor const &logit_lengths, torch::Tensor const &target_lengths, torch::Tensor const &durations, torch::Tensor const &out, @@ -173,7 +173,7 @@ void check_lattice_args(torch::Tensor const &blank_lp, } TORCH_CHECK(out.scalar_type() == torch::kFloat64, "alphas/betas must be float64"); - for (auto const *x : {&source_lengths, &target_lengths, &durations}) { + for (auto const *x : {&logit_lengths, &target_lengths, &durations}) { TORCH_CHECK(x->device() == blank_lp.device(), "all tensors must be on the same device"); TORCH_CHECK(x->scalar_type() == torch::kInt32, @@ -186,7 +186,7 @@ void check_lattice_args(torch::Tensor const &blank_lp, "lattice tensors must have shape (batch, T, U+1)"); TORCH_CHECK(dur_lp.numel() == nodes * durations.numel(), "dur_lp must have shape (batch, T, U+1, num_durations)"); - TORCH_CHECK(source_lengths.numel() == B && target_lengths.numel() == B, + TORCH_CHECK(logit_lengths.numel() == B && target_lengths.numel() == B, "lengths must have shape (batch,)"); if (log_ll != nullptr) { TORCH_CHECK(log_ll->device() == blank_lp.device() && @@ -204,12 +204,12 @@ void check_lattice_args(torch::Tensor const &blank_lp, 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 &logit_lengths, torch::Tensor const &target_lengths, torch::Tensor const &durations, torch::Tensor &alphas, torch::Tensor &log_ll) { using namespace tdt_loss; - check_lattice_args(blank_lp, label_lp, dur_lp, source_lengths, + check_lattice_args(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations, alphas, &log_ll); const int64_t B = blank_lp.size(0); if (B == 0) return; @@ -218,7 +218,7 @@ void tdt_loss_fwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); tdt_alpha_kernel<<>>( blank_lp.data_ptr(), label_lp.data_ptr(), - dur_lp.data_ptr(), source_lengths.data_ptr(), + dur_lp.data_ptr(), logit_lengths.data_ptr(), target_lengths.data_ptr(), durations.data_ptr(), blank_lp.size(1), blank_lp.size(2), durations.numel(), alphas.data_ptr(), log_ll.data_ptr()); @@ -227,11 +227,11 @@ void tdt_loss_fwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, 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 &logit_lengths, torch::Tensor const &target_lengths, torch::Tensor const &durations, torch::Tensor &betas) { using namespace tdt_loss; - check_lattice_args(blank_lp, label_lp, dur_lp, source_lengths, + check_lattice_args(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations, betas, nullptr); const int64_t B = blank_lp.size(0); if (B == 0) return; @@ -240,7 +240,7 @@ void tdt_loss_bwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); tdt_beta_kernel<<>>( blank_lp.data_ptr(), label_lp.data_ptr(), - dur_lp.data_ptr(), source_lengths.data_ptr(), + dur_lp.data_ptr(), logit_lengths.data_ptr(), target_lengths.data_ptr(), durations.data_ptr(), blank_lp.size(1), blank_lp.size(2), durations.numel(), betas.data_ptr()); diff --git a/tdt-loss/tdt_loss/logprobs.cu b/tdt-loss/tdt_loss/logprobs.cu index bb54d7b44..db3b9e0c9 100644 --- a/tdt-loss/tdt_loss/logprobs.cu +++ b/tdt-loss/tdt_loss/logprobs.cu @@ -23,7 +23,7 @@ __global__ void __launch_bounds__(kMaxRowThreads) tdt_logprobs_fwd_kernel( int64_t tok_su, const scalar_t *__restrict__ duration_logits, int64_t dur_sb, int64_t dur_st, int64_t dur_su, const int *__restrict__ targets, int64_t targets_stride, - const int *__restrict__ source_lengths, + const int *__restrict__ logit_lengths, const int *__restrict__ target_lengths, int max_T, int max_U, int V, int D, int blank_id, float sigma, float *__restrict__ blank_lp, float *__restrict__ label_lp, float *__restrict__ dur_lp, @@ -37,7 +37,7 @@ __global__ void __launch_bounds__(kMaxRowThreads) tdt_logprobs_fwd_kernel( const int b = row / (static_cast(max_U) * max_T); int T, U; - sample_lengths(source_lengths, target_lengths, b, max_T, max_U, T, U); + sample_lengths(logit_lengths, target_lengths, b, max_T, max_U, T, U); if (t >= T || u > U) { if (threadIdx.x == 0) { @@ -101,7 +101,7 @@ template __global__ void __launch_bounds__(kMaxRowThreads) tdt_logits_grad_kernel( const scalar_t *__restrict__ token_logits, int64_t tok_sb, int64_t tok_st, int64_t tok_su, const int *__restrict__ targets, int64_t targets_stride, - const int *__restrict__ source_lengths, + const int *__restrict__ logit_lengths, const int *__restrict__ target_lengths, const int *__restrict__ durations, const float *__restrict__ blank_lp, const float *__restrict__ label_lp, const float *__restrict__ dur_lp, const float *__restrict__ token_lse, @@ -121,7 +121,7 @@ __global__ void __launch_bounds__(kMaxRowThreads) tdt_logits_grad_kernel( const int b = row / (static_cast(max_U) * max_T); int T, U; - sample_lengths(source_lengths, target_lengths, b, max_T, max_U, T, U); + sample_lengths(logit_lengths, target_lengths, b, max_T, max_U, T, U); scalar_t *gx = grad_token_logits + row * V; scalar_t *gd = grad_duration_logits + row * D; @@ -256,7 +256,7 @@ void check_float_out(torch::Tensor const &x, char const *name, 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 &logit_lengths, torch::Tensor const &target_lengths, int64_t blank_id, double sigma, torch::Tensor &blank_lp, torch::Tensor &label_lp, torch::Tensor &dur_lp, @@ -275,9 +275,9 @@ void tdt_logprobs_fwd(torch::Tensor const &token_logits, targets.size(1) >= max_U - 1, "targets must have shape (batch, U)"); check_int_tensor(targets, "targets", token_logits); - check_int_tensor(source_lengths, "source_lengths", token_logits); + check_int_tensor(logit_lengths, "logit_lengths", token_logits); check_int_tensor(target_lengths, "target_lengths", token_logits); - TORCH_CHECK(source_lengths.numel() == B && target_lengths.numel() == B, + TORCH_CHECK(logit_lengths.numel() == B && target_lengths.numel() == B, "lengths must have shape (batch,)"); for (auto *out : {&blank_lp, &label_lp, &token_lse}) { check_float_out(*out, "log-prob output", token_logits); @@ -306,7 +306,7 @@ void tdt_logprobs_fwd(torch::Tensor const &token_logits, duration_logits.data_ptr(), duration_logits.stride(0), duration_logits.stride(1), duration_logits.stride(2), targets.data_ptr(), targets.stride(0), - source_lengths.data_ptr(), target_lengths.data_ptr(), + logit_lengths.data_ptr(), target_lengths.data_ptr(), max_T, max_U, V, D, blank_id, static_cast(sigma), blank_lp.data_ptr(), label_lp.data_ptr(), dur_lp.data_ptr(), token_lse.data_ptr()); @@ -316,7 +316,7 @@ void tdt_logprobs_fwd(torch::Tensor const &token_logits, void tdt_logits_grad( torch::Tensor const &token_logits, torch::Tensor const &targets, - torch::Tensor const &source_lengths, torch::Tensor const &target_lengths, + torch::Tensor const &logit_lengths, torch::Tensor const &target_lengths, torch::Tensor const &durations, torch::Tensor const &blank_lp, torch::Tensor const &label_lp, torch::Tensor const &dur_lp, torch::Tensor const &token_lse, torch::Tensor const &alphas, @@ -339,7 +339,7 @@ void tdt_logits_grad( " durations are supported"); TORCH_CHECK(blank_id >= 0 && blank_id < V, "blank_id out of range"); check_int_tensor(targets, "targets", token_logits); - check_int_tensor(source_lengths, "source_lengths", token_logits); + check_int_tensor(logit_lengths, "logit_lengths", token_logits); check_int_tensor(target_lengths, "target_lengths", token_logits); check_int_tensor(durations, "durations", token_logits); for (auto *in : {&blank_lp, &label_lp, &token_lse}) { @@ -389,7 +389,7 @@ void tdt_logits_grad( token_logits.data_ptr(), token_logits.stride(0), token_logits.stride(1), token_logits.stride(2), targets.data_ptr(), targets.stride(0), - source_lengths.data_ptr(), target_lengths.data_ptr(), + logit_lengths.data_ptr(), target_lengths.data_ptr(), durations.data_ptr(), blank_lp.data_ptr(), label_lp.data_ptr(), dur_lp.data_ptr(), token_lse.data_ptr(), alphas.data_ptr(), diff --git a/tdt-loss/tests/test_tdt_loss.py b/tdt-loss/tests/test_tdt_loss.py index a0af3d3a6..7c060ff5c 100644 --- a/tdt-loss/tests/test_tdt_loss.py +++ b/tdt-loss/tests/test_tdt_loss.py @@ -336,3 +336,60 @@ def test_layer(): reduction="none", ) torch.testing.assert_close(tdt_loss_kernel.layers.TDTLoss()(**kwargs), tdt_loss_kernel.tdt_loss(**kwargs)) + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_torch_compile(dtype): + """The loss and its gradients are the same with `torch.compile(fullgraph=True)`.""" + if dtype == torch.bfloat16 and torch.cuda.get_device_capability() < (8, 0): + pytest.skip("bfloat16 requires compute capability >= 8.0") + V = 64 + token_logits, _, targets, logit_lengths, target_lengths = make_inputs(3, 20, 6, V, DURATIONS, dtype=dtype, joint=True) + + def loss_fn(logits, targets, logit_lengths, target_lengths): + return tdt_loss_kernel.tdt_loss( + logits[..., :V], logits[..., V:], targets, logit_lengths, target_lengths, 0, DURATIONS, sigma=0.05 + ) + + torch._dynamo.reset() + results = [] + for fn in (loss_fn, torch.compile(loss_fn, fullgraph=True)): + logits = token_logits._base.detach().clone().requires_grad_(True) + loss = fn(logits, targets, logit_lengths, target_lengths) + loss.backward() + results.append((loss.detach(), logits.grad)) + + torch.testing.assert_close(results[1][0], results[0][0]) + torch.testing.assert_close(results[1][1], results[0][1]) + + +@pytest.mark.kernels_ci +def test_opcheck(): + """The schemas and fake implementations of the custom ops are consistent (`torch.library.opcheck`).""" + ops = tdt_loss_kernel.loss.ops + B, T, U, V = 2, 8, 4, 16 + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs(B, T, U, V, DURATIONS) + durations = torch.tensor(DURATIONS, device=DEVICE, dtype=torch.int32) + f32 = dict(device=DEVICE, dtype=torch.float32) + f64 = dict(device=DEVICE, dtype=torch.float64) + shape = (B, T, U + 1) + + blank_lp, label_lp, token_lse = (torch.empty(shape, **f32) for _ in range(3)) + dur_lp = torch.empty(*shape, len(DURATIONS), **f32) + args = (token_logits, duration_logits, targets, logit_lengths, target_lengths, 0, 0.0) + args += (blank_lp, label_lp, dur_lp, token_lse) + torch.library.opcheck(ops.tdt_logprobs_fwd, args) + ops.tdt_logprobs_fwd(*args) + + lattice = (blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations) + alphas, betas, log_ll = torch.empty(shape, **f64), torch.empty(shape, **f64), torch.empty(B, **f64) + torch.library.opcheck(ops.tdt_loss_fwd, lattice + (alphas, log_ll)) + ops.tdt_loss_fwd(*lattice, alphas, log_ll) + torch.library.opcheck(ops.tdt_loss_bwd, lattice + (betas,)) + ops.tdt_loss_bwd(*lattice, betas) + + args = (token_logits, targets, logit_lengths, target_lengths, durations, blank_lp, label_lp, dur_lp, token_lse) + args += (alphas, betas, log_ll, torch.ones(B, **f32), 0) + args += (torch.empty_like(token_logits), torch.empty_like(duration_logits)) + torch.library.opcheck(ops.tdt_logits_grad, args) diff --git a/tdt-loss/torch-ext/tdt_loss/layers.py b/tdt-loss/torch-ext/tdt_loss/layers.py index fcede711c..fe75c1eec 100644 --- a/tdt-loss/torch-ext/tdt_loss/layers.py +++ b/tdt-loss/torch-ext/tdt_loss/layers.py @@ -13,7 +13,7 @@ class TDTLoss(nn.Module): """ has_backward = True - can_torch_compile = False + can_torch_compile = True def forward( self, diff --git a/tdt-loss/torch-ext/tdt_loss/loss.py b/tdt-loss/torch-ext/tdt_loss/loss.py index 01f646501..3beb979b1 100644 --- a/tdt-loss/torch-ext/tdt_loss/loss.py +++ b/tdt-loss/torch-ext/tdt_loss/loss.py @@ -4,7 +4,7 @@ import torch -from ._ops import ops +from ._ops import add_op_namespace_prefix, ops _REDUCTIONS = ("mean_volume", "mean_batch", "mean", "sum", "none") @@ -13,13 +13,15 @@ def _durations_tensor(durations: Union[Sequence[int], torch.Tensor], device: torch.device) -> torch.Tensor: - """Validated int32 tensor of durations on `device`, cached since the durations are fixed for a model.""" + """Validated int32 tensor of durations on `device`, cached in eager mode since they are fixed for a model.""" if isinstance(durations, torch.Tensor): durations = durations.tolist() key = (tuple(int(d) for d in durations), device) + if any(d < 0 for d in key[0]): + raise ValueError(f"Durations must be non-negative, got {list(key[0])}.") + if torch.compiler.is_compiling(): + return torch.tensor(key[0], device=device, dtype=torch.int32) if key not in _DURATIONS_CACHE: - if any(d < 0 for d in key[0]): - raise ValueError(f"Durations must be non-negative, got {list(key[0])}.") _DURATIONS_CACHE[key] = torch.tensor(key[0], device=device, dtype=torch.int32) return _DURATIONS_CACHE[key] @@ -30,6 +32,57 @@ def _last_dim_contiguous(x: torch.Tensor) -> torch.Tensor: return x if x.stride(-1) == 1 else x.contiguous() +# The ops only write to preallocated outputs, so their fake (meta) implementations have nothing to do. They let +# `torch.compile` trace through the ops. +@torch.library.register_fake(add_op_namespace_prefix("tdt_logprobs_fwd")) +def _( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + blank_id, + sigma, + blank_lp, + label_lp, + dur_lp, + token_lse, +): + return None + + +@torch.library.register_fake(add_op_namespace_prefix("tdt_loss_fwd")) +def _(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations, alphas, log_ll): + return None + + +@torch.library.register_fake(add_op_namespace_prefix("tdt_loss_bwd")) +def _(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations, betas): + return None + + +@torch.library.register_fake(add_op_namespace_prefix("tdt_logits_grad")) +def _( + token_logits, + targets, + logit_lengths, + target_lengths, + durations, + blank_lp, + label_lp, + dur_lp, + token_lse, + alphas, + betas, + log_ll, + grad_loss, + blank_id, + grad_token_logits, + grad_duration_logits, +): + return None + + class _TDTLossFunction(torch.autograd.Function): """Per-sample TDT loss (negative log-likelihood) with a CUDA forward and backward.""" diff --git a/tdt-loss/torch-ext/torch_binding.cpp b/tdt-loss/torch-ext/torch_binding.cpp index b3d91f420..fef1bf113 100644 --- a/tdt-loss/torch-ext/torch_binding.cpp +++ b/tdt-loss/torch-ext/torch_binding.cpp @@ -6,26 +6,26 @@ 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, " + "Tensor targets, Tensor logit_lengths, Tensor target_lengths, " "int blank_id, float sigma, Tensor! blank_lp, Tensor! label_lp, " "Tensor! dur_lp, Tensor! token_lse) -> ()"); 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 logit_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 logit_lengths, Tensor target_lengths, Tensor durations, " "Tensor! betas) -> ()"); ops.impl("tdt_loss_bwd", torch::kCUDA, &tdt_loss_bwd); ops.def( "tdt_logits_grad(Tensor token_logits, Tensor targets, " - "Tensor source_lengths, Tensor target_lengths, Tensor durations, " + "Tensor logit_lengths, Tensor target_lengths, Tensor durations, " "Tensor blank_lp, Tensor label_lp, Tensor dur_lp, Tensor token_lse, " "Tensor alphas, Tensor betas, Tensor log_ll, Tensor grad_loss, " "int blank_id, Tensor! grad_token_logits, " diff --git a/tdt-loss/torch-ext/torch_binding.h b/tdt-loss/torch-ext/torch_binding.h index c1337b4ec..97b903dc0 100644 --- a/tdt-loss/torch-ext/torch_binding.h +++ b/tdt-loss/torch-ext/torch_binding.h @@ -8,7 +8,7 @@ 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 &logit_lengths, torch::Tensor const &target_lengths, int64_t blank_id, double sigma, torch::Tensor &blank_lp, torch::Tensor &label_lp, torch::Tensor &dur_lp, @@ -17,7 +17,7 @@ void tdt_logprobs_fwd(torch::Tensor const &token_logits, // Forward recursion over the lattice (alphas) and per-sample log-likelihood. 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 &logit_lengths, torch::Tensor const &target_lengths, torch::Tensor const &durations, torch::Tensor &alphas, torch::Tensor &log_ll); @@ -25,7 +25,7 @@ void tdt_loss_fwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, // Backward recursion over the lattice (betas). 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 &logit_lengths, torch::Tensor const &target_lengths, torch::Tensor const &durations, torch::Tensor &betas); @@ -33,7 +33,7 @@ void tdt_loss_bwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, // and duration logits. void tdt_logits_grad( torch::Tensor const &token_logits, torch::Tensor const &targets, - torch::Tensor const &source_lengths, torch::Tensor const &target_lengths, + torch::Tensor const &logit_lengths, torch::Tensor const &target_lengths, torch::Tensor const &durations, torch::Tensor const &blank_lp, torch::Tensor const &label_lp, torch::Tensor const &dur_lp, torch::Tensor const &token_lse, torch::Tensor const &alphas, From ab47fc95266dd258c19496830724c15eb42585b9 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 5 Oct 2026 17:33:19 +0200 Subject: [PATCH 14/17] tdt-loss: add maintainer, drop README (the card is used for the Hub) --- .github/kernel-maintainers.json | 2 +- tdt-loss/README.md | 84 --------------------------------- 2 files changed, 1 insertion(+), 85 deletions(-) delete mode 100644 tdt-loss/README.md diff --git a/.github/kernel-maintainers.json b/.github/kernel-maintainers.json index 923a929b6..3215df6f9 100644 --- a/.github/kernel-maintainers.json +++ b/.github/kernel-maintainers.json @@ -51,7 +51,7 @@ "scattermoe": ["U08Q42K8MGW"], "sgl-flash-attn3": [], "sonic-moe": ["U08Q42K8MGW", "U054KJK89C2"], - "tdt-loss": [], + "tdt-loss": ["U07U53UJECT"], "tinygrad-rms": [], "topk": ["U0586KAH2E9"], "trimul-gpumode": [], diff --git a/tdt-loss/README.md b/tdt-loss/README.md deleted file mode 100644 index e3cbb6314..000000000 --- a/tdt-loss/README.md +++ /dev/null @@ -1,84 +0,0 @@ ---- -tags: -- kernel ---- -CUDA kernels for the TDT (Token-and-Duration Transducer, [arXiv:2304.06795](https://arxiv.org/abs/2304.06795)) loss, -used to train [Parakeet TDT](https://huggingface.co/docs/transformers/model_doc/parakeet) models in 🤗 Transformers. - -The forward and backward passes are computed with dedicated kernels: - -- a log-softmax over the vocabulary for every lattice node (one thread block per node, online softmax), -- the alpha / beta recursions over the anti-diagonals of the lattice (one thread block per sample), -- a fused gradient w.r.t. the token and duration logits. - -Logits can be float32, float16 or bfloat16 and only need a contiguous last dimension, so slices of a joint -`(..., vocab_size + num_durations)` output are used without a copy. The lattice recursions run in float64. - -## Usage - -```python -# /// script -# dependencies = [ -# "torch", -# "kernels" -# ] -# /// -import torch -from kernels import get_kernel - -tdt = get_kernel("kernels-community/tdt-loss", version=1) - -batch_size, max_frames, max_labels, vocab_size = 2, 100, 20, 1024 -durations = [0, 1, 2, 3, 4] -blank_id = vocab_size # blank is the last token -device = "cuda" - -# Joint network output: vocabulary (with blank) followed by the durations. -logits = torch.randn( - batch_size, max_frames, max_labels + 1, vocab_size + 1 + len(durations), device=device, requires_grad=True -) -targets = torch.randint(0, vocab_size, (batch_size, max_labels), device=device) -logit_lengths = torch.tensor([max_frames, max_frames - 10], device=device) -target_lengths = torch.tensor([max_labels, max_labels - 5], device=device) - -loss = tdt.tdt_loss( - logits[..., : vocab_size + 1], - logits[..., vocab_size + 1 :], - targets, - logit_lengths, - target_lengths, - blank_id, - durations, - sigma=0.0, - reduction="mean", # or "mean_volume", "mean_batch", "sum", "none" -) -loss.backward() -``` - -Samples without a valid alignment get an infinite loss and a zero gradient. - -The loss supports `torch.compile` (including `fullgraph=True`). The `layers.TDTLoss` layer has the same signature -as `transformers.loss.loss_tdt.tdt_loss`, so 🤗 Transformers swaps it in for Parakeet TDT models loaded with -`use_kernels=True`. - -## Benchmarks - -Forward + backward of the loss, with the logits given as slices of a joint output, compared with the vectorized -PyTorch implementation from `transformers.loss.loss_tdt`. Run `benchmarks/benchmark.py` with `kernels benchmark`. - -| batch, frames, labels, vocab | dtype | GPU | PyTorch | kernel | speedup | -|---|---|---|---:|---:|---:| -| 8, 100, 20, 1025 | float32 | L4 | 536 ms | 4.3 ms | 125x | -| 4, 200, 40, 8193 | float32 | L4 | 1252 ms | 53 ms | 24x | -| 4, 200, 40, 8193 | bfloat16 | L4 | 1242 ms | 39 ms | 32x | -| 8, 100, 20, 1025 | float32 | A100 | 624 ms | 1.5 ms | 410x | -| 4, 200, 40, 8193 | float32 | A100 | 1259 ms | 9.5 ms | 133x | -| 4, 200, 40, 8193 | bfloat16 | A100 | 1267 ms | 8.6 ms | 148x | -| 8, 100, 20, 1025 | float32 | H200 | 488 ms | 1.0 ms | 480x | -| 4, 200, 40, 8193 | float32 | H200 | 993 ms | 5.0 ms | 199x | -| 4, 200, 40, 8193 | bfloat16 | H200 | 896 ms | 4.9 ms | 183x | - -Fine-tuning [nvidia/parakeet-tdt-0.6b-v3](https://huggingface.co/nvidia/parakeet-tdt-0.6b-v3) on LibriSpeech -(batch size 8, bfloat16 autocast, 300 steps, same data order and seed) on an H200 gives the same loss curve as the -PyTorch loss (identical first-step loss, median per-step difference of 0.0015), with a full training step going from -1305 ms to 144 ms (9x) and peak memory from 33.2 GiB to 21.5 GiB. From febfe211d0e931829d9ce897ef8f32e383de0a02 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 5 Oct 2026 17:33:19 +0200 Subject: [PATCH 15/17] tdt-loss: move to the Torch stable ABI Use torch::stable::Tensor, STD_TORCH_CHECK, THO dispatch, the stable DeviceGuard and the AOTI shim for the current stream, and register the ops with STABLE_TORCH_LIBRARY. The CUDA kernels are unchanged. --- tdt-loss/build.toml | 4 + tdt-loss/tdt_loss/common.cuh | 24 ++- tdt-loss/tdt_loss/lattice.cu | 206 +++++++++--------- tdt-loss/tdt_loss/logprobs.cu | 298 ++++++++++++++------------- tdt-loss/tdt_loss/stable.h | 54 +++++ tdt-loss/torch-ext/torch_binding.cpp | 51 +++-- tdt-loss/torch-ext/torch_binding.h | 69 ++++--- 7 files changed, 386 insertions(+), 320 deletions(-) create mode 100644 tdt-loss/tdt_loss/stable.h diff --git a/tdt-loss/build.toml b/tdt-loss/build.toml index f83977d85..a135eebc2 100644 --- a/tdt-loss/build.toml +++ b/tdt-loss/build.toml @@ -13,6 +13,9 @@ src = [ "torch-ext/torch_binding.h", ] +[torch.stable-abi] +cuda = "2.10" + [kernel.tdt_loss] backend = "cuda" depends = ["torch"] @@ -20,4 +23,5 @@ src = [ "tdt_loss/common.cuh", "tdt_loss/lattice.cu", "tdt_loss/logprobs.cu", + "tdt_loss/stable.h", ] diff --git a/tdt-loss/tdt_loss/common.cuh b/tdt-loss/tdt_loss/common.cuh index de02e7ca1..f3dbb4b19 100644 --- a/tdt-loss/tdt_loss/common.cuh +++ b/tdt-loss/tdt_loss/common.cuh @@ -1,7 +1,7 @@ #pragma once -#include #include +#include #include #include @@ -23,8 +23,10 @@ __device__ __forceinline__ float neg_inf() { return -INFINITY; } // float64 to keep large log-likelihoods exact, but the correction term only // depends on |a - b| and is computed in float32. __device__ __forceinline__ double log_add(double a, double b) { - if (a == -INFINITY) return b; - if (b == -INFINITY) return a; + if (a == -INFINITY) + return b; + if (b == -INFINITY) + return a; return fmax(a, b) + log1pf(expf(-fabsf(static_cast(a - b)))); } @@ -35,7 +37,8 @@ template __device__ __forceinline__ double log_sum_exp(ForEach for_each) { double m = -INFINITY; for_each([&](double v) { m = fmax(m, v); }); - if (m == -INFINITY) return m; + if (m == -INFINITY) + return m; float s = 0.f; for_each([&](double v) { s += expf(static_cast(v - m)); }); return m + logf(s); @@ -56,9 +59,8 @@ __device__ __forceinline__ void merge_max_sum(float &m, float &s, float m_other, // Block-wide reduction of (max, sum) pairs. The result is valid in all threads. // `shm_m` and `shm_s` must hold at least blockDim.x / kWarpSize floats. -__device__ __forceinline__ void block_reduce_max_sum(float &m, float &s, - float *shm_m, - float *shm_s) { +__device__ __forceinline__ void +block_reduce_max_sum(float &m, float &s, float *shm_m, float *shm_s) { const int lane = threadIdx.x % kWarpSize; const int warp = threadIdx.x / kWarpSize; const int num_warps = (blockDim.x + kWarpSize - 1) / kWarpSize; @@ -99,9 +101,9 @@ __device__ __forceinline__ void block_reduce_max_sum(float &m, float &s, // they are also clamped so that no out-of-bounds access happens when asserts // are compiled out. __device__ __forceinline__ void sample_lengths(const int *logit_lengths, - const int *target_lengths, - int b, int max_T, int max_U, - int &T, int &U) { + const int *target_lengths, int b, + int max_T, int max_U, int &T, + int &U) { T = logit_lengths[b]; U = target_lengths[b]; CUDA_KERNEL_ASSERT(T >= 0 && T <= max_T && "logit_lengths out of range"); @@ -110,4 +112,4 @@ __device__ __forceinline__ void sample_lengths(const int *logit_lengths, U = min(max(U, 0), max_U - 1); } -} // namespace tdt_loss +} // namespace tdt_loss diff --git a/tdt-loss/tdt_loss/lattice.cu b/tdt-loss/tdt_loss/lattice.cu index 27e1a94e0..d77c86462 100644 --- a/tdt-loss/tdt_loss/lattice.cu +++ b/tdt-loss/tdt_loss/lattice.cu @@ -14,11 +14,8 @@ // cost, and the gradient combines alpha + beta - log_likelihood, which cancels // catastrophically in float32 when the log-likelihood is large (long targets). -#include -#include -#include - #include "common.cuh" +#include "stable.h" namespace tdt_loss { @@ -32,15 +29,12 @@ __device__ __forceinline__ void load_durations(const int *durations, int D, __syncthreads(); } -__global__ void __launch_bounds__(kMaxLatticeThreads) - tdt_alpha_kernel(const float *__restrict__ blank_lp, - const float *__restrict__ label_lp, - const float *__restrict__ dur_lp, - const int *__restrict__ logit_lengths, - const int *__restrict__ target_lengths, - const int *__restrict__ durations, int max_T, - int max_U, int D, double *__restrict__ alphas, - double *__restrict__ log_ll) { +__global__ void __launch_bounds__(kMaxLatticeThreads) tdt_alpha_kernel( + const float *__restrict__ blank_lp, const float *__restrict__ label_lp, + const float *__restrict__ dur_lp, const int *__restrict__ logit_lengths, + const int *__restrict__ target_lengths, const int *__restrict__ durations, + int max_T, int max_U, int D, double *__restrict__ alphas, + double *__restrict__ log_ll) { __shared__ int shm_durations[kMaxDurations]; const int b = blockIdx.x; @@ -59,23 +53,22 @@ __global__ void __launch_bounds__(kMaxLatticeThreads) const int u_max = min(n, U); for (int u = u_min + threadIdx.x; u <= u_max; u += blockDim.x) { const int t = n - u; - const double a = - n == 0 ? 0.0 : log_sum_exp([&](auto &&f) { - for (int i = 0; i < D; ++i) { - const int d = shm_durations[i]; - const int t_prev = t - d; - if (t_prev < 0) continue; - if (d > 0) { - const int64_t src = static_cast(t_prev) * max_U + u; - f(alpha[src] + blank[src] + dur[src * D + i]); - } - if (u > 0) { - const int64_t src = - static_cast(t_prev) * max_U + u - 1; - f(alpha[src] + label[src] + dur[src * D + i]); - } - } - }); + const double a = n == 0 ? 0.0 : log_sum_exp([&](auto &&f) { + for (int i = 0; i < D; ++i) { + const int d = shm_durations[i]; + const int t_prev = t - d; + if (t_prev < 0) + continue; + if (d > 0) { + const int64_t src = static_cast(t_prev) * max_U + u; + f(alpha[src] + blank[src] + dur[src * D + i]); + } + if (u > 0) { + const int64_t src = static_cast(t_prev) * max_U + u - 1; + f(alpha[src] + label[src] + dur[src * D + i]); + } + } + }); alpha[static_cast(t) * max_U + u] = a; } __syncthreads(); @@ -86,7 +79,8 @@ __global__ void __launch_bounds__(kMaxLatticeThreads) for (int i = 0; i < D; ++i) { const int d = shm_durations[i]; const int t_prev = T - d; - if (d == 0 || t_prev < 0) continue; + if (d == 0 || t_prev < 0) + continue; const int64_t src = static_cast(t_prev) * max_U + U; ll = log_add(ll, alpha[src] + blank[src] + dur[src * D + i]); } @@ -94,14 +88,11 @@ __global__ void __launch_bounds__(kMaxLatticeThreads) } } -__global__ void __launch_bounds__(kMaxLatticeThreads) - tdt_beta_kernel(const float *__restrict__ blank_lp, - const float *__restrict__ label_lp, - const float *__restrict__ dur_lp, - const int *__restrict__ logit_lengths, - const int *__restrict__ target_lengths, - const int *__restrict__ durations, int max_T, int max_U, - int D, double *__restrict__ betas) { +__global__ void __launch_bounds__(kMaxLatticeThreads) tdt_beta_kernel( + const float *__restrict__ blank_lp, const float *__restrict__ label_lp, + const float *__restrict__ dur_lp, const int *__restrict__ logit_lengths, + const int *__restrict__ target_lengths, const int *__restrict__ durations, + int max_T, int max_U, int D, double *__restrict__ betas) { __shared__ int shm_durations[kMaxDurations]; const int b = blockIdx.x; @@ -131,7 +122,7 @@ __global__ void __launch_bounds__(kMaxLatticeThreads) const int64_t dst = static_cast(t_next) * max_U + u; f(beta[dst] + blank[node] + lp_dur); } else if (t_next == T && u == U) { - f(blank[node] + lp_dur); // terminal arc + f(blank[node] + lp_dur); // terminal arc } } if (u < U && t_next < T) { @@ -148,101 +139,100 @@ __global__ void __launch_bounds__(kMaxLatticeThreads) int lattice_block_size(int64_t max_U) { int threads = 32; - while (threads < kMaxLatticeThreads && threads < max_U) threads *= 2; + while (threads < kMaxLatticeThreads && threads < max_U) + threads *= 2; return threads; } -void check_lattice_args(torch::Tensor const &blank_lp, - torch::Tensor const &label_lp, - torch::Tensor const &dur_lp, - torch::Tensor const &logit_lengths, - torch::Tensor const &target_lengths, - torch::Tensor const &durations, - torch::Tensor const &out, - torch::Tensor const *log_ll) { - TORCH_CHECK(blank_lp.is_cuda(), "blank_lp must be a CUDA tensor"); - TORCH_CHECK(blank_lp.dim() == 3, "blank_lp must have shape (batch, T, U+1)"); +void check_lattice_args(Tensor const &blank_lp, Tensor const &label_lp, + Tensor const &dur_lp, Tensor const &logit_lengths, + Tensor const &target_lengths, Tensor const &durations, + Tensor const &out, Tensor const *log_ll) { + STD_TORCH_CHECK(blank_lp.is_cuda(), "blank_lp must be a CUDA tensor"); + STD_TORCH_CHECK(blank_lp.dim() == 3, + "blank_lp must have shape (batch, T, U+1)"); for (auto const *x : {&blank_lp, &label_lp, &dur_lp, &out}) { - TORCH_CHECK(x->device() == blank_lp.device(), - "all tensors must be on the same device"); - TORCH_CHECK(x->is_contiguous(), "lattice tensors must be contiguous"); + STD_TORCH_CHECK(same_device(*x, blank_lp), + "all tensors must be on the same device"); + STD_TORCH_CHECK(x->is_contiguous(), "lattice tensors must be contiguous"); } for (auto const *x : {&blank_lp, &label_lp, &dur_lp}) { - TORCH_CHECK(x->scalar_type() == torch::kFloat32, - "log-prob tensors must be float32"); + STD_TORCH_CHECK(x->scalar_type() == ScalarType::Float, + "log-prob tensors must be float32"); } - TORCH_CHECK(out.scalar_type() == torch::kFloat64, - "alphas/betas must be float64"); + STD_TORCH_CHECK(out.scalar_type() == ScalarType::Double, + "alphas/betas must be float64"); for (auto const *x : {&logit_lengths, &target_lengths, &durations}) { - TORCH_CHECK(x->device() == blank_lp.device(), - "all tensors must be on the same device"); - TORCH_CHECK(x->scalar_type() == torch::kInt32, - "lengths and durations must be int32"); - TORCH_CHECK(x->is_contiguous(), "lengths and durations must be contiguous"); + STD_TORCH_CHECK(same_device(*x, blank_lp), + "all tensors must be on the same device"); + STD_TORCH_CHECK(x->scalar_type() == ScalarType::Int, + "lengths and durations must be int32"); + STD_TORCH_CHECK(x->is_contiguous(), + "lengths and durations must be contiguous"); } const int64_t B = blank_lp.size(0); const int64_t nodes = blank_lp.numel(); - TORCH_CHECK(label_lp.numel() == nodes && out.numel() == nodes, - "lattice tensors must have shape (batch, T, U+1)"); - TORCH_CHECK(dur_lp.numel() == nodes * durations.numel(), - "dur_lp must have shape (batch, T, U+1, num_durations)"); - TORCH_CHECK(logit_lengths.numel() == B && target_lengths.numel() == B, - "lengths must have shape (batch,)"); + STD_TORCH_CHECK(label_lp.numel() == nodes && out.numel() == nodes, + "lattice tensors must have shape (batch, T, U+1)"); + STD_TORCH_CHECK(dur_lp.numel() == nodes * durations.numel(), + "dur_lp must have shape (batch, T, U+1, num_durations)"); + STD_TORCH_CHECK(logit_lengths.numel() == B && target_lengths.numel() == B, + "lengths must have shape (batch,)"); if (log_ll != nullptr) { - TORCH_CHECK(log_ll->device() == blank_lp.device() && - log_ll->scalar_type() == torch::kFloat64 && - log_ll->is_contiguous() && log_ll->numel() == B, - "log_ll must be a contiguous float64 tensor of shape (batch,)"); + STD_TORCH_CHECK( + same_device(*log_ll, blank_lp) && + log_ll->scalar_type() == ScalarType::Double && + log_ll->is_contiguous() && log_ll->numel() == B, + "log_ll must be a contiguous float64 tensor of shape (batch,)"); } - TORCH_CHECK(durations.numel() <= kMaxDurations, "at most ", kMaxDurations, - " durations are supported"); + STD_TORCH_CHECK(durations.numel() <= kMaxDurations, "at most ", kMaxDurations, + " durations are supported"); } -} // namespace +} // namespace -} // namespace tdt_loss +} // namespace tdt_loss -void tdt_loss_fwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, - torch::Tensor const &dur_lp, - torch::Tensor const &logit_lengths, - torch::Tensor const &target_lengths, - torch::Tensor const &durations, torch::Tensor &alphas, - torch::Tensor &log_ll) { +void tdt_loss_fwd(Tensor const &blank_lp, Tensor const &label_lp, + Tensor const &dur_lp, Tensor const &logit_lengths, + Tensor const &target_lengths, Tensor const &durations, + Tensor &alphas, Tensor &log_ll) { using namespace tdt_loss; - check_lattice_args(blank_lp, label_lp, dur_lp, logit_lengths, - target_lengths, durations, alphas, &log_ll); + check_lattice_args(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, + durations, alphas, &log_ll); const int64_t B = blank_lp.size(0); - if (B == 0) return; + if (B == 0) + return; - const at::cuda::OptionalCUDAGuard device_guard(device_of(blank_lp)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + blank_lp.get_device_index()); + const cudaStream_t stream = current_stream(blank_lp); tdt_alpha_kernel<<>>( - blank_lp.data_ptr(), label_lp.data_ptr(), - dur_lp.data_ptr(), logit_lengths.data_ptr(), - target_lengths.data_ptr(), durations.data_ptr(), + ptr(blank_lp), ptr(label_lp), ptr(dur_lp), + ptr(logit_lengths), ptr(target_lengths), ptr(durations), blank_lp.size(1), blank_lp.size(2), durations.numel(), - alphas.data_ptr(), log_ll.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + ptr(alphas), ptr(log_ll)); + STD_CUDA_KERNEL_LAUNCH_CHECK(); } -void tdt_loss_bwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, - torch::Tensor const &dur_lp, - torch::Tensor const &logit_lengths, - torch::Tensor const &target_lengths, - torch::Tensor const &durations, torch::Tensor &betas) { +void tdt_loss_bwd(Tensor const &blank_lp, Tensor const &label_lp, + Tensor const &dur_lp, Tensor const &logit_lengths, + Tensor const &target_lengths, Tensor const &durations, + Tensor &betas) { using namespace tdt_loss; - check_lattice_args(blank_lp, label_lp, dur_lp, logit_lengths, - target_lengths, durations, betas, nullptr); + check_lattice_args(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, + durations, betas, nullptr); const int64_t B = blank_lp.size(0); - if (B == 0) return; + if (B == 0) + return; - const at::cuda::OptionalCUDAGuard device_guard(device_of(blank_lp)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + blank_lp.get_device_index()); + const cudaStream_t stream = current_stream(blank_lp); tdt_beta_kernel<<>>( - blank_lp.data_ptr(), label_lp.data_ptr(), - dur_lp.data_ptr(), logit_lengths.data_ptr(), - target_lengths.data_ptr(), durations.data_ptr(), + ptr(blank_lp), ptr(label_lp), ptr(dur_lp), + ptr(logit_lengths), ptr(target_lengths), ptr(durations), blank_lp.size(1), blank_lp.size(2), durations.numel(), - betas.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + ptr(betas)); + STD_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/tdt-loss/tdt_loss/logprobs.cu b/tdt-loss/tdt_loss/logprobs.cu index db3b9e0c9..e5f0cfcb5 100644 --- a/tdt-loss/tdt_loss/logprobs.cu +++ b/tdt-loss/tdt_loss/logprobs.cu @@ -5,13 +5,10 @@ // vocabulary dimension. Nodes outside of a sample's lattice (t >= T or u > U) // never read the logits. -#include -#include -#include - #include #include "common.cuh" +#include "stable.h" namespace tdt_loss { @@ -84,12 +81,13 @@ __global__ void __launch_bounds__(kMaxRowThreads) tdt_logprobs_fwd_kernel( // The duration head is tiny (a handful of classes), a serial // log-softmax is cheaper than another block reduction. - const scalar_t *y = - duration_logits + b * dur_sb + t * dur_st + u * dur_su; + const scalar_t *y = duration_logits + b * dur_sb + t * dur_st + u * dur_su; float dm = neg_inf(); - for (int i = 0; i < D; ++i) dm = fmaxf(dm, static_cast(y[i])); + for (int i = 0; i < D; ++i) + dm = fmaxf(dm, static_cast(y[i])); float ds = 0.f; - for (int i = 0; i < D; ++i) ds += expf(static_cast(y[i]) - dm); + for (int i = 0; i < D; ++i) + ds += expf(static_cast(y[i]) - dm); const float dlse = dm + logf(ds); for (int i = 0; i < D; ++i) { dur_lp[row * D + i] = static_cast(y[i]) - dlse; @@ -130,8 +128,10 @@ __global__ void __launch_bounds__(kMaxRowThreads) tdt_logits_grad_kernel( const double ll = log_ll[b]; // Infeasible samples (infinite loss) get a zero gradient instead of NaNs. if (t >= T || u > U || !isfinite(ll)) { - for (int v = threadIdx.x; v < V; v += blockDim.x) gx[v] = zero; - for (int i = threadIdx.x; i < D; i += blockDim.x) gd[i] = zero; + for (int v = threadIdx.x; v < V; v += blockDim.x) + gx[v] = zero; + for (int i = threadIdx.x; i < D; i += blockDim.x) + gd[i] = zero; return; } @@ -151,10 +151,11 @@ __global__ void __launch_bounds__(kMaxRowThreads) tdt_logits_grad_kernel( if (t_next < T) { beta_next = betas[sample + static_cast(t_next) * max_U + u]; } else if (t_next == T && u == U) { - beta_next = 0.0; // terminal arc + beta_next = 0.0; // terminal arc } if (beta_next != -INFINITY) { - p_blank = expf(static_cast(alpha + blank + lp_dur + beta_next - ll)); + p_blank = + expf(static_cast(alpha + blank + lp_dur + beta_next - ll)); } } @@ -162,7 +163,8 @@ __global__ void __launch_bounds__(kMaxRowThreads) tdt_logits_grad_kernel( if (u < U && t_next < T) { const double beta_next = betas[sample + static_cast(t_next) * max_U + u + 1]; - p_label = expf(static_cast(alpha + label + lp_dur + beta_next - ll)); + p_label = + expf(static_cast(alpha + label + lp_dur + beta_next - ll)); } shm_blank[i] = p_blank; @@ -190,7 +192,8 @@ __global__ void __launch_bounds__(kMaxRowThreads) tdt_logits_grad_kernel( } if (g_blank == 0.f && g_label == 0.f) { - for (int v = threadIdx.x; v < V; v += blockDim.x) gx[v] = zero; + for (int v = threadIdx.x; v < V; v += blockDim.x) + gx[v] = zero; return; } @@ -199,68 +202,68 @@ __global__ void __launch_bounds__(kMaxRowThreads) tdt_logits_grad_kernel( const int target = u < U ? targets[b * targets_stride + u] : -1; for (int v = threadIdx.x; v < V; v += blockDim.x) { float g = -expf(static_cast(x[v]) - lse) * g_total; - if (v == blank_id) g += g_blank; - if (v == target) g += g_label; + if (v == blank_id) + g += g_blank; + if (v == target) + g += g_label; gx[v] = static_cast(g); } } int row_block_size(int64_t V) { int threads = 32; - while (threads < kMaxRowThreads && threads < V) threads *= 2; + while (threads < kMaxRowThreads && threads < V) + threads *= 2; return threads; } -void check_logits(torch::Tensor const &token_logits, - torch::Tensor const &duration_logits) { - TORCH_CHECK(token_logits.is_cuda(), "token_logits must be a CUDA tensor"); - TORCH_CHECK(token_logits.dim() == 4, - "token_logits must have shape (batch, T, U+1, vocab_size+1)"); - TORCH_CHECK(duration_logits.dim() == 4, - "duration_logits must have shape (batch, T, U+1, num_durations)"); - TORCH_CHECK(duration_logits.device() == token_logits.device(), - "token_logits and duration_logits must be on the same device"); - TORCH_CHECK(duration_logits.scalar_type() == token_logits.scalar_type(), - "token_logits and duration_logits must have the same dtype"); +void check_logits(Tensor const &token_logits, Tensor const &duration_logits) { + STD_TORCH_CHECK(token_logits.is_cuda(), "token_logits must be a CUDA tensor"); + STD_TORCH_CHECK(token_logits.dim() == 4, + "token_logits must have shape (batch, T, U+1, vocab_size+1)"); + STD_TORCH_CHECK( + duration_logits.dim() == 4, + "duration_logits must have shape (batch, T, U+1, num_durations)"); + STD_TORCH_CHECK( + same_device(duration_logits, token_logits), + "token_logits and duration_logits must be on the same device"); + STD_TORCH_CHECK(duration_logits.scalar_type() == token_logits.scalar_type(), + "token_logits and duration_logits must have the same dtype"); for (int i = 0; i < 3; ++i) { - TORCH_CHECK(duration_logits.size(i) == token_logits.size(i), - "token_logits and duration_logits must agree on (batch, T, " - "U+1)"); + STD_TORCH_CHECK(duration_logits.size(i) == token_logits.size(i), + "token_logits and duration_logits must agree on (batch, T, " + "U+1)"); } - TORCH_CHECK(token_logits.stride(3) == 1 && duration_logits.stride(3) == 1, - "the last dimension of the logits must be contiguous"); - TORCH_CHECK(duration_logits.size(3) <= kMaxDurations, - "at most ", kMaxDurations, " durations are supported"); + STD_TORCH_CHECK(token_logits.stride(3) == 1 && duration_logits.stride(3) == 1, + "the last dimension of the logits must be contiguous"); + STD_TORCH_CHECK(duration_logits.size(3) <= kMaxDurations, "at most ", + kMaxDurations, " durations are supported"); } -void check_int_tensor(torch::Tensor const &x, char const *name, - torch::Tensor const &like) { - TORCH_CHECK(x.device() == like.device(), name, - " must be on the same device as the logits"); - TORCH_CHECK(x.scalar_type() == torch::kInt32, name, " must be int32"); - TORCH_CHECK(x.is_contiguous(), name, " must be contiguous"); +void check_int_tensor(Tensor const &x, char const *name, Tensor const &like) { + STD_TORCH_CHECK(same_device(x, like), name, + " must be on the same device as the logits"); + STD_TORCH_CHECK(x.scalar_type() == ScalarType::Int, name, " must be int32"); + STD_TORCH_CHECK(x.is_contiguous(), name, " must be contiguous"); } -void check_float_out(torch::Tensor const &x, char const *name, - torch::Tensor const &like) { - TORCH_CHECK(x.device() == like.device(), name, - " must be on the same device as the logits"); - TORCH_CHECK(x.scalar_type() == torch::kFloat32, name, " must be float32"); - TORCH_CHECK(x.is_contiguous(), name, " must be contiguous"); +void check_float_out(Tensor const &x, char const *name, Tensor const &like) { + STD_TORCH_CHECK(same_device(x, like), name, + " must be on the same device as the logits"); + STD_TORCH_CHECK(x.scalar_type() == ScalarType::Float, name, + " must be float32"); + STD_TORCH_CHECK(x.is_contiguous(), name, " must be contiguous"); } -} // namespace +} // namespace -} // namespace tdt_loss +} // namespace tdt_loss -void tdt_logprobs_fwd(torch::Tensor const &token_logits, - torch::Tensor const &duration_logits, - torch::Tensor const &targets, - torch::Tensor const &logit_lengths, - torch::Tensor const &target_lengths, int64_t blank_id, - double sigma, torch::Tensor &blank_lp, - torch::Tensor &label_lp, torch::Tensor &dur_lp, - torch::Tensor &token_lse) { +void tdt_logprobs_fwd(Tensor const &token_logits, Tensor const &duration_logits, + Tensor const &targets, Tensor const &logit_lengths, + Tensor const &target_lengths, int64_t blank_id, + double sigma, Tensor &blank_lp, Tensor &label_lp, + Tensor &dur_lp, Tensor &token_lse) { using namespace tdt_loss; check_logits(token_logits, duration_logits); @@ -270,133 +273,136 @@ void tdt_logprobs_fwd(torch::Tensor const &token_logits, const int64_t V = token_logits.size(3); const int64_t D = duration_logits.size(3); - TORCH_CHECK(blank_id >= 0 && blank_id < V, "blank_id out of range"); - TORCH_CHECK(targets.dim() == 2 && targets.size(0) == B && - targets.size(1) >= max_U - 1, - "targets must have shape (batch, U)"); + STD_TORCH_CHECK(blank_id >= 0 && blank_id < V, "blank_id out of range"); + STD_TORCH_CHECK(targets.dim() == 2 && targets.size(0) == B && + targets.size(1) >= max_U - 1, + "targets must have shape (batch, U)"); check_int_tensor(targets, "targets", token_logits); check_int_tensor(logit_lengths, "logit_lengths", token_logits); check_int_tensor(target_lengths, "target_lengths", token_logits); - TORCH_CHECK(logit_lengths.numel() == B && target_lengths.numel() == B, - "lengths must have shape (batch,)"); + STD_TORCH_CHECK(logit_lengths.numel() == B && target_lengths.numel() == B, + "lengths must have shape (batch,)"); for (auto *out : {&blank_lp, &label_lp, &token_lse}) { check_float_out(*out, "log-prob output", token_logits); - TORCH_CHECK(out->numel() == B * max_T * max_U, - "log-prob outputs must have shape (batch, T, U+1)"); + STD_TORCH_CHECK(out->numel() == B * max_T * max_U, + "log-prob outputs must have shape (batch, T, U+1)"); } check_float_out(dur_lp, "dur_lp", token_logits); - TORCH_CHECK(dur_lp.numel() == B * max_T * max_U * D, - "dur_lp must have shape (batch, T, U+1, num_durations)"); + STD_TORCH_CHECK(dur_lp.numel() == B * max_T * max_U * D, + "dur_lp must have shape (batch, T, U+1, num_durations)"); const int64_t rows = B * max_T * max_U; - if (rows == 0) return; - TORCH_CHECK(rows <= std::numeric_limits::max(), - "batch * T * (U+1) is too large"); + if (rows == 0) + return; + STD_TORCH_CHECK(rows <= std::numeric_limits::max(), + "batch * T * (U+1) is too large"); - const at::cuda::OptionalCUDAGuard device_guard(device_of(token_logits)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + token_logits.get_device_index()); + const cudaStream_t stream = current_stream(token_logits); const int threads = row_block_size(V); - AT_DISPATCH_FLOATING_TYPES_AND2( - at::kHalf, at::kBFloat16, token_logits.scalar_type(), "tdt_logprobs_fwd", - [&] { - tdt_logprobs_fwd_kernel<<(rows), threads, 0, stream>>>( - token_logits.data_ptr(), token_logits.stride(0), - token_logits.stride(1), token_logits.stride(2), - duration_logits.data_ptr(), duration_logits.stride(0), - duration_logits.stride(1), duration_logits.stride(2), - targets.data_ptr(), targets.stride(0), - logit_lengths.data_ptr(), target_lengths.data_ptr(), - max_T, max_U, V, D, blank_id, static_cast(sigma), - blank_lp.data_ptr(), label_lp.data_ptr(), - dur_lp.data_ptr(), token_lse.data_ptr()); + TDT_DISPATCH_FLOATING_TYPES( + token_logits.scalar_type(), "tdt_logprobs_fwd", [&] { + tdt_logprobs_fwd_kernel + <<(rows), threads, 0, stream>>>( + ptr(token_logits), token_logits.stride(0), + token_logits.stride(1), token_logits.stride(2), + ptr(duration_logits), duration_logits.stride(0), + duration_logits.stride(1), duration_logits.stride(2), + ptr(targets), targets.stride(0), ptr(logit_lengths), + ptr(target_lengths), max_T, max_U, V, D, blank_id, + static_cast(sigma), ptr(blank_lp), + ptr(label_lp), ptr(dur_lp), + ptr(token_lse)); }); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + STD_CUDA_KERNEL_LAUNCH_CHECK(); } -void tdt_logits_grad( - torch::Tensor const &token_logits, torch::Tensor const &targets, - torch::Tensor const &logit_lengths, torch::Tensor const &target_lengths, - torch::Tensor const &durations, torch::Tensor const &blank_lp, - torch::Tensor const &label_lp, torch::Tensor const &dur_lp, - torch::Tensor const &token_lse, torch::Tensor const &alphas, - torch::Tensor const &betas, torch::Tensor const &log_ll, - torch::Tensor const &grad_loss, int64_t blank_id, - torch::Tensor &grad_token_logits, torch::Tensor &grad_duration_logits) { +void tdt_logits_grad(Tensor const &token_logits, Tensor const &targets, + Tensor const &logit_lengths, Tensor const &target_lengths, + Tensor const &durations, Tensor const &blank_lp, + Tensor const &label_lp, Tensor const &dur_lp, + Tensor const &token_lse, Tensor const &alphas, + Tensor const &betas, Tensor const &log_ll, + Tensor const &grad_loss, int64_t blank_id, + Tensor &grad_token_logits, Tensor &grad_duration_logits) { using namespace tdt_loss; - TORCH_CHECK(token_logits.is_cuda() && token_logits.dim() == 4, - "token_logits must be a 4D CUDA tensor"); - TORCH_CHECK(token_logits.stride(3) == 1, - "the last dimension of token_logits must be contiguous"); + STD_TORCH_CHECK(token_logits.is_cuda() && token_logits.dim() == 4, + "token_logits must be a 4D CUDA tensor"); + STD_TORCH_CHECK(token_logits.stride(3) == 1, + "the last dimension of token_logits must be contiguous"); const int64_t B = token_logits.size(0); const int64_t max_T = token_logits.size(1); const int64_t max_U = token_logits.size(2); const int64_t V = token_logits.size(3); const int64_t D = durations.numel(); - TORCH_CHECK(D <= kMaxDurations, "at most ", kMaxDurations, - " durations are supported"); - TORCH_CHECK(blank_id >= 0 && blank_id < V, "blank_id out of range"); + STD_TORCH_CHECK(D <= kMaxDurations, "at most ", kMaxDurations, + " durations are supported"); + STD_TORCH_CHECK(blank_id >= 0 && blank_id < V, "blank_id out of range"); check_int_tensor(targets, "targets", token_logits); check_int_tensor(logit_lengths, "logit_lengths", token_logits); check_int_tensor(target_lengths, "target_lengths", token_logits); check_int_tensor(durations, "durations", token_logits); for (auto *in : {&blank_lp, &label_lp, &token_lse}) { check_float_out(*in, "log-prob tensor", token_logits); - TORCH_CHECK(in->numel() == B * max_T * max_U, - "log-prob tensors must have shape (batch, T, U+1)"); + STD_TORCH_CHECK(in->numel() == B * max_T * max_U, + "log-prob tensors must have shape (batch, T, U+1)"); } for (auto *in : {&alphas, &betas}) { - TORCH_CHECK(in->device() == token_logits.device() && - in->scalar_type() == torch::kFloat64 && - in->is_contiguous() && in->numel() == B * max_T * max_U, - "alphas and betas must be contiguous float64 tensors of shape " - "(batch, T, U+1)"); + STD_TORCH_CHECK( + same_device(*in, token_logits) && + in->scalar_type() == ScalarType::Double && in->is_contiguous() && + in->numel() == B * max_T * max_U, + "alphas and betas must be contiguous float64 tensors of shape " + "(batch, T, U+1)"); } check_float_out(dur_lp, "dur_lp", token_logits); - TORCH_CHECK(log_ll.device() == token_logits.device() && - log_ll.scalar_type() == torch::kFloat64 && - log_ll.is_contiguous(), - "log_ll must be a contiguous float64 tensor"); + STD_TORCH_CHECK(same_device(log_ll, token_logits) && + log_ll.scalar_type() == ScalarType::Double && + log_ll.is_contiguous(), + "log_ll must be a contiguous float64 tensor"); check_float_out(grad_loss, "grad_loss", token_logits); - TORCH_CHECK(log_ll.numel() == B && grad_loss.numel() == B, - "log_ll and grad_loss must have shape (batch,)"); - TORCH_CHECK(grad_token_logits.is_contiguous() && - grad_token_logits.sizes() == token_logits.sizes() && - grad_token_logits.scalar_type() == token_logits.scalar_type(), - "grad_token_logits must be contiguous and match token_logits"); - TORCH_CHECK(grad_duration_logits.is_contiguous() && - grad_duration_logits.numel() == B * max_T * max_U * D && - grad_duration_logits.scalar_type() == - token_logits.scalar_type(), - "grad_duration_logits must be contiguous and match " - "duration_logits"); + STD_TORCH_CHECK(log_ll.numel() == B && grad_loss.numel() == B, + "log_ll and grad_loss must have shape (batch,)"); + STD_TORCH_CHECK( + grad_token_logits.is_contiguous() && + grad_token_logits.sizes() == token_logits.sizes() && + grad_token_logits.scalar_type() == token_logits.scalar_type(), + "grad_token_logits must be contiguous and match token_logits"); + STD_TORCH_CHECK(grad_duration_logits.is_contiguous() && + grad_duration_logits.numel() == B * max_T * max_U * D && + grad_duration_logits.scalar_type() == + token_logits.scalar_type(), + "grad_duration_logits must be contiguous and match " + "duration_logits"); const int64_t rows = B * max_T * max_U; - if (rows == 0) return; - TORCH_CHECK(rows <= std::numeric_limits::max(), - "batch * T * (U+1) is too large"); + if (rows == 0) + return; + STD_TORCH_CHECK(rows <= std::numeric_limits::max(), + "batch * T * (U+1) is too large"); - const at::cuda::OptionalCUDAGuard device_guard(device_of(token_logits)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + token_logits.get_device_index()); + const cudaStream_t stream = current_stream(token_logits); const int threads = row_block_size(V); - AT_DISPATCH_FLOATING_TYPES_AND2( - at::kHalf, at::kBFloat16, token_logits.scalar_type(), "tdt_logits_grad", - [&] { - tdt_logits_grad_kernel<<(rows), threads, 0, stream>>>( - token_logits.data_ptr(), token_logits.stride(0), - token_logits.stride(1), token_logits.stride(2), - targets.data_ptr(), targets.stride(0), - logit_lengths.data_ptr(), target_lengths.data_ptr(), - durations.data_ptr(), blank_lp.data_ptr(), - label_lp.data_ptr(), dur_lp.data_ptr(), - token_lse.data_ptr(), alphas.data_ptr(), - betas.data_ptr(), log_ll.data_ptr(), - grad_loss.data_ptr(), max_T, max_U, V, D, blank_id, - grad_token_logits.data_ptr(), - grad_duration_logits.data_ptr()); + TDT_DISPATCH_FLOATING_TYPES( + token_logits.scalar_type(), "tdt_logits_grad", [&] { + tdt_logits_grad_kernel + <<(rows), threads, 0, stream>>>( + ptr(token_logits), token_logits.stride(0), + token_logits.stride(1), token_logits.stride(2), + ptr(targets), targets.stride(0), ptr(logit_lengths), + ptr(target_lengths), ptr(durations), + ptr(blank_lp), ptr(label_lp), ptr(dur_lp), + ptr(token_lse), ptr(alphas), ptr(betas), + ptr(log_ll), ptr(grad_loss), max_T, max_U, V, D, + blank_id, ptr(grad_token_logits), + ptr(grad_duration_logits)); }); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + STD_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/tdt-loss/tdt_loss/stable.h b/tdt-loss/tdt_loss/stable.h new file mode 100644 index 000000000..f9e6baf6e --- /dev/null +++ b/tdt-loss/tdt_loss/stable.h @@ -0,0 +1,54 @@ +#pragma once + +// Helpers for the host code on top of the Torch stable ABI. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// The shim's stream accessor is guarded by USE_CUDA, so declare it here. +extern "C" AOTITorchError +aoti_torch_get_current_cuda_stream(int32_t device_index, void **ret_stream); + +using torch::headeronly::ScalarType; +using torch::stable::Tensor; + +// Dispatches on the logits dtype, with the type available as `scalar_t`. +#define TDT_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH( \ + TYPE, NAME, \ + THO_PRIVATE_CASE_TYPE_USING_HINT(ScalarType::Float, scalar_t, \ + __VA_ARGS__) \ + THO_PRIVATE_CASE_TYPE_USING_HINT(ScalarType::Half, scalar_t, \ + __VA_ARGS__) \ + THO_PRIVATE_CASE_TYPE_USING_HINT(ScalarType::BFloat16, scalar_t, \ + __VA_ARGS__)) + +namespace tdt_loss { + +template T *ptr(Tensor const &x) { + return static_cast(x.data_ptr()); +} + +inline bool same_device(Tensor const &a, Tensor const &b) { + return a.is_cuda() == b.is_cuda() && + a.get_device_index() == b.get_device_index(); +} + +inline cudaStream_t current_stream(Tensor const &x) { + void *stream = nullptr; + TORCH_ERROR_CODE_CHECK( + aoti_torch_get_current_cuda_stream(x.get_device_index(), &stream)); + return static_cast(stream); +} + +} // namespace tdt_loss diff --git a/tdt-loss/torch-ext/torch_binding.cpp b/tdt-loss/torch-ext/torch_binding.cpp index fef1bf113..f41039291 100644 --- a/tdt-loss/torch-ext/torch_binding.cpp +++ b/tdt-loss/torch-ext/torch_binding.cpp @@ -1,36 +1,35 @@ -#include +#include #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 logit_lengths, Tensor target_lengths, " - "int blank_id, float sigma, Tensor! blank_lp, Tensor! label_lp, " - "Tensor! dur_lp, Tensor! token_lse) -> ()"); - ops.impl("tdt_logprobs_fwd", torch::kCUDA, &tdt_logprobs_fwd); +STABLE_TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { + ops.def("tdt_logprobs_fwd(Tensor token_logits, Tensor duration_logits, " + "Tensor targets, Tensor logit_lengths, Tensor target_lengths, " + "int blank_id, float sigma, Tensor! blank_lp, Tensor! label_lp, " + "Tensor! dur_lp, Tensor! token_lse) -> ()"); - ops.def( - "tdt_loss_fwd(Tensor blank_lp, Tensor label_lp, Tensor dur_lp, " - "Tensor logit_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_fwd(Tensor blank_lp, Tensor label_lp, Tensor dur_lp, " + "Tensor logit_lengths, Tensor target_lengths, Tensor durations, " + "Tensor! alphas, Tensor! log_ll) -> ()"); - ops.def( - "tdt_loss_bwd(Tensor blank_lp, Tensor label_lp, Tensor dur_lp, " - "Tensor logit_lengths, Tensor target_lengths, Tensor durations, " - "Tensor! betas) -> ()"); - ops.impl("tdt_loss_bwd", torch::kCUDA, &tdt_loss_bwd); + ops.def("tdt_loss_bwd(Tensor blank_lp, Tensor label_lp, Tensor dur_lp, " + "Tensor logit_lengths, Tensor target_lengths, Tensor durations, " + "Tensor! betas) -> ()"); - ops.def( - "tdt_logits_grad(Tensor token_logits, Tensor targets, " - "Tensor logit_lengths, Tensor target_lengths, Tensor durations, " - "Tensor blank_lp, Tensor label_lp, Tensor dur_lp, Tensor token_lse, " - "Tensor alphas, Tensor betas, Tensor log_ll, Tensor grad_loss, " - "int blank_id, Tensor! grad_token_logits, " - "Tensor! grad_duration_logits) -> ()"); - ops.impl("tdt_logits_grad", torch::kCUDA, &tdt_logits_grad); + ops.def("tdt_logits_grad(Tensor token_logits, Tensor targets, " + "Tensor logit_lengths, Tensor target_lengths, Tensor durations, " + "Tensor blank_lp, Tensor label_lp, Tensor dur_lp, Tensor token_lse, " + "Tensor alphas, Tensor betas, Tensor log_ll, Tensor grad_loss, " + "int blank_id, Tensor! grad_token_logits, " + "Tensor! grad_duration_logits) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, ops) { + ops.impl("tdt_logprobs_fwd", TORCH_BOX(&tdt_logprobs_fwd)); + ops.impl("tdt_loss_fwd", TORCH_BOX(&tdt_loss_fwd)); + ops.impl("tdt_loss_bwd", TORCH_BOX(&tdt_loss_bwd)); + ops.impl("tdt_logits_grad", TORCH_BOX(&tdt_logits_grad)); } REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/tdt-loss/torch-ext/torch_binding.h b/tdt-loss/torch-ext/torch_binding.h index 97b903dc0..3acc065cc 100644 --- a/tdt-loss/torch-ext/torch_binding.h +++ b/tdt-loss/torch-ext/torch_binding.h @@ -1,42 +1,53 @@ #pragma once -#include +#include // Log-softmax gather over the vocabulary and duration heads. Writes the blank // and label log-probabilities of every lattice node, the duration // log-probabilities and the token log-normalizer (reused by the backward). -void tdt_logprobs_fwd(torch::Tensor const &token_logits, - torch::Tensor const &duration_logits, - torch::Tensor const &targets, - torch::Tensor const &logit_lengths, - torch::Tensor const &target_lengths, int64_t blank_id, - double sigma, torch::Tensor &blank_lp, - torch::Tensor &label_lp, torch::Tensor &dur_lp, - torch::Tensor &token_lse); +void tdt_logprobs_fwd(torch::stable::Tensor const &token_logits, + torch::stable::Tensor const &duration_logits, + torch::stable::Tensor const &targets, + torch::stable::Tensor const &logit_lengths, + torch::stable::Tensor const &target_lengths, + int64_t blank_id, double sigma, + torch::stable::Tensor &blank_lp, + torch::stable::Tensor &label_lp, + torch::stable::Tensor &dur_lp, + torch::stable::Tensor &token_lse); // Forward recursion over the lattice (alphas) and per-sample log-likelihood. -void tdt_loss_fwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, - torch::Tensor const &dur_lp, - torch::Tensor const &logit_lengths, - torch::Tensor const &target_lengths, - torch::Tensor const &durations, torch::Tensor &alphas, - torch::Tensor &log_ll); +void tdt_loss_fwd(torch::stable::Tensor const &blank_lp, + torch::stable::Tensor const &label_lp, + torch::stable::Tensor const &dur_lp, + torch::stable::Tensor const &logit_lengths, + torch::stable::Tensor const &target_lengths, + torch::stable::Tensor const &durations, + torch::stable::Tensor &alphas, torch::stable::Tensor &log_ll); // Backward recursion over the lattice (betas). -void tdt_loss_bwd(torch::Tensor const &blank_lp, torch::Tensor const &label_lp, - torch::Tensor const &dur_lp, - torch::Tensor const &logit_lengths, - torch::Tensor const &target_lengths, - torch::Tensor const &durations, torch::Tensor &betas); +void tdt_loss_bwd(torch::stable::Tensor const &blank_lp, + torch::stable::Tensor const &label_lp, + torch::stable::Tensor const &dur_lp, + torch::stable::Tensor const &logit_lengths, + torch::stable::Tensor const &target_lengths, + torch::stable::Tensor const &durations, + torch::stable::Tensor &betas); // Gradient of the per-sample losses (scaled by grad_loss) w.r.t. the token // and duration logits. -void tdt_logits_grad( - torch::Tensor const &token_logits, torch::Tensor const &targets, - torch::Tensor const &logit_lengths, torch::Tensor const &target_lengths, - torch::Tensor const &durations, torch::Tensor const &blank_lp, - torch::Tensor const &label_lp, torch::Tensor const &dur_lp, - torch::Tensor const &token_lse, torch::Tensor const &alphas, - torch::Tensor const &betas, torch::Tensor const &log_ll, - torch::Tensor const &grad_loss, int64_t blank_id, - torch::Tensor &grad_token_logits, torch::Tensor &grad_duration_logits); +void tdt_logits_grad(torch::stable::Tensor const &token_logits, + torch::stable::Tensor const &targets, + torch::stable::Tensor const &logit_lengths, + torch::stable::Tensor const &target_lengths, + torch::stable::Tensor const &durations, + torch::stable::Tensor const &blank_lp, + torch::stable::Tensor const &label_lp, + torch::stable::Tensor const &dur_lp, + torch::stable::Tensor const &token_lse, + torch::stable::Tensor const &alphas, + torch::stable::Tensor const &betas, + torch::stable::Tensor const &log_ll, + torch::stable::Tensor const &grad_loss, int64_t blank_id, + torch::stable::Tensor &grad_token_logits, + torch::stable::Tensor &grad_duration_logits); From 474a55f7115008cac1ef6cd8e5ba3c8fd2197310 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Mon, 5 Oct 2026 18:03:25 +0200 Subject: [PATCH 16/17] tdt-loss: stable ABI fixes for Torch 2.10 --- tdt-loss/tdt_loss/lattice.cu | 4 ++-- tdt-loss/tdt_loss/logprobs.cu | 6 +++--- tdt-loss/tdt_loss/stable.h | 19 +++++++++++++++++++ 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/tdt-loss/tdt_loss/lattice.cu b/tdt-loss/tdt_loss/lattice.cu index d77c86462..76608df73 100644 --- a/tdt-loss/tdt_loss/lattice.cu +++ b/tdt-loss/tdt_loss/lattice.cu @@ -212,7 +212,7 @@ void tdt_loss_fwd(Tensor const &blank_lp, Tensor const &label_lp, ptr(logit_lengths), ptr(target_lengths), ptr(durations), blank_lp.size(1), blank_lp.size(2), durations.numel(), ptr(alphas), ptr(log_ll)); - STD_CUDA_KERNEL_LAUNCH_CHECK(); + TDT_CUDA_KERNEL_LAUNCH_CHECK(); } void tdt_loss_bwd(Tensor const &blank_lp, Tensor const &label_lp, @@ -234,5 +234,5 @@ void tdt_loss_bwd(Tensor const &blank_lp, Tensor const &label_lp, ptr(logit_lengths), ptr(target_lengths), ptr(durations), blank_lp.size(1), blank_lp.size(2), durations.numel(), ptr(betas)); - STD_CUDA_KERNEL_LAUNCH_CHECK(); + TDT_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/tdt-loss/tdt_loss/logprobs.cu b/tdt-loss/tdt_loss/logprobs.cu index e5f0cfcb5..b74a64e34 100644 --- a/tdt-loss/tdt_loss/logprobs.cu +++ b/tdt-loss/tdt_loss/logprobs.cu @@ -316,7 +316,7 @@ void tdt_logprobs_fwd(Tensor const &token_logits, Tensor const &duration_logits, ptr(label_lp), ptr(dur_lp), ptr(token_lse)); }); - STD_CUDA_KERNEL_LAUNCH_CHECK(); + TDT_CUDA_KERNEL_LAUNCH_CHECK(); } void tdt_logits_grad(Tensor const &token_logits, Tensor const &targets, @@ -369,7 +369,7 @@ void tdt_logits_grad(Tensor const &token_logits, Tensor const &targets, "log_ll and grad_loss must have shape (batch,)"); STD_TORCH_CHECK( grad_token_logits.is_contiguous() && - grad_token_logits.sizes() == token_logits.sizes() && + same_sizes(grad_token_logits, token_logits) && grad_token_logits.scalar_type() == token_logits.scalar_type(), "grad_token_logits must be contiguous and match token_logits"); STD_TORCH_CHECK(grad_duration_logits.is_contiguous() && @@ -404,5 +404,5 @@ void tdt_logits_grad(Tensor const &token_logits, Tensor const &targets, blank_id, ptr(grad_token_logits), ptr(grad_duration_logits)); }); - STD_CUDA_KERNEL_LAUNCH_CHECK(); + TDT_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/tdt-loss/tdt_loss/stable.h b/tdt-loss/tdt_loss/stable.h index f9e6baf6e..20b6b65cf 100644 --- a/tdt-loss/tdt_loss/stable.h +++ b/tdt-loss/tdt_loss/stable.h @@ -33,6 +33,15 @@ using torch::stable::Tensor; THO_PRIVATE_CASE_TYPE_USING_HINT(ScalarType::BFloat16, scalar_t, \ __VA_ARGS__)) +// Like STD_CUDA_KERNEL_LAUNCH_CHECK, which needs shim functions that are only +// declared with USE_CUDA. +#define TDT_CUDA_KERNEL_LAUNCH_CHECK() \ + do { \ + const cudaError_t err = cudaGetLastError(); \ + STD_TORCH_CHECK(err == cudaSuccess, \ + "CUDA error: ", cudaGetErrorString(err)); \ + } while (0) + namespace tdt_loss { template T *ptr(Tensor const &x) { @@ -44,6 +53,16 @@ inline bool same_device(Tensor const &a, Tensor const &b) { a.get_device_index() == b.get_device_index(); } +inline bool same_sizes(Tensor const &a, Tensor const &b) { + if (a.dim() != b.dim()) + return false; + for (int64_t i = 0; i < a.dim(); ++i) { + if (a.size(i) != b.size(i)) + return false; + } + return true; +} + inline cudaStream_t current_stream(Tensor const &x) { void *stream = nullptr; TORCH_ERROR_CODE_CHECK( From 1eefa6f976267c980fb91022fd524a1d9ee6c924 Mon Sep 17 00:00:00 2001 From: Deep-unlearning Date: Tue, 6 Oct 2026 10:53:21 +0200 Subject: [PATCH 17/17] tdt-loss: use build.toml edition 5, bump kernel-builder --- tdt-loss/build.toml | 1 + tdt-loss/flake.lock | 20 ++++++++++---------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/tdt-loss/build.toml b/tdt-loss/build.toml index a135eebc2..7028ad9cd 100644 --- a/tdt-loss/build.toml +++ b/tdt-loss/build.toml @@ -2,6 +2,7 @@ name = "tdt-loss" license = "Apache-2.0" version = 1 +edition = 5 backends = ["cuda"] [general.hub] diff --git a/tdt-loss/flake.lock b/tdt-loss/flake.lock index 35a669620..f80d937fa 100644 --- a/tdt-loss/flake.lock +++ b/tdt-loss/flake.lock @@ -41,11 +41,11 @@ "rust-overlay": "rust-overlay" }, "locked": { - "lastModified": 1788544112, - "narHash": "sha256-aqwWRYCri6ZTSMQNRE0EQZRvP45nubvcW62pVhSeoyg=", + "lastModified": 1790766967, + "narHash": "sha256-P4QLsZU4S/ZOaxIgc3KXTdaZ3uHJ+R7ko0ACKZG31vA=", "owner": "huggingface", "repo": "kernels", - "rev": "a7f0afdb29a6a3372b1d47180cc0c182454c5e3b", + "rev": "a56da98e235a3f52b4b766e4d93fac6a9e599e3d", "type": "github" }, "original": { @@ -56,17 +56,17 @@ }, "nixpkgs": { "locked": { - "lastModified": 1783284758, - "narHash": "sha256-tiQ8/qi8I45OOaBBYlVbXoAVkeQzvvTQOv5I45rMw5o=", + "lastModified": 1789524109, + "narHash": "sha256-GZDvT/tbSVQB2KjI86l7VCM5D/vCrJo2sE3PCaaD0oQ=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d", + "rev": "be3dab5b1a3c41ed2db86093b9f4e048305630b7", "type": "github" }, "original": { "owner": "NixOS", + "ref": "nixos-unstable-small", "repo": "nixpkgs", - "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d", "type": "github" } }, @@ -83,11 +83,11 @@ ] }, "locked": { - "lastModified": 1783320166, - "narHash": "sha256-l7C/OsjcnWDOk2K3ssj+SBduwL67LashjBqis9+t468=", + "lastModified": 1789543870, + "narHash": "sha256-45ikFmtfgBZgMePIm4GevJ9VjDSWO3PGQgOuuwHDpes=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "20ee15370c9256669d66968b89ee20a4b0a4e673", + "rev": "35ca0490d13a3d38c4602d0eb9600a30fa63a367", "type": "github" }, "original": {