liger-kernel 0.5.3__py3-none-any.whl → 0.5.5__py3-none-any.whl

liger_kernel/chunked_loss/kto_loss.py

@@ -7,10 +7,10 @@ from liger_kernel.chunked_loss.fused_linear_unpaired_preference import LigerFuse
 class LigerFusedLinearKTOFunction(LigerFusedLinearUnpairedPreferenceBase):
     @staticmethod
     def preference_loss_fn(
-        average_log_prob_chunk,
+        log_prob_chunk,
         preference_labels_chunk,
         full_target,
-        ref_average_log_prob_chunk=None,
+        ref_log_prob_chunk=None,
         beta=0.1,
         kl=None,
     ):
@@ -43,30 +43,34 @@ class LigerFusedLinearKTOFunction(LigerFusedLinearUnpairedPreferenceBase):
         3. Maintain reasonable distance from the reference model
 
         Args:
-            chosen_logps: Log probabilities of chosen tokens (batch_size,)
-            rejected_logps: Log probabilities of rejected tokens (batch_size,)
+            log_prob_chunk: Log probabilities for the chunk (batch_size,)
+            preference_labels_chunk: Preference labels for the chunk (batch_size,)
             full_target: Non chunked full target tensor
-            ref_chosen_logps: Reference log probs of chosen tokens (batch_size,)
-            ref_rejected_logps: Reference log probs of rejected tokens (batch_size,)
-            beta: Weight for the direct preference loss
+            ref_log_prob_chunk: Reference log probs for the chunk (batch_size,)
+            beta: Weight for the KTO loss
             kl: KL divergence between the policy model and the reference model for the chosen responses. Shape: (batch_size,)
         Returns:
-            Tuple of (loss, chosen_rewards, rejected_rewards):
             - loss: The KTO loss value
-            - chosen_rewards: Reward signals for chosen responses (detached)
-            - rejected_rewards: Reward signals for rejected responses (detached)
         """
-        logratios_chunk = average_log_prob_chunk - ref_average_log_prob_chunk
+        if ref_log_prob_chunk is not None:
+            logratios_chunk = log_prob_chunk - ref_log_prob_chunk
+        else:
+            logratios_chunk = log_prob_chunk
         multiplier_chunk = torch.where(preference_labels_chunk, 1, -1)
         if kl is not None:
             losses = 1 - F.sigmoid(beta * (logratios_chunk - kl) * multiplier_chunk)
         else:
             losses = 1 - F.sigmoid(beta * logratios_chunk * multiplier_chunk)
 
-        return losses.sum() / (full_target.shape[0])
+        rewards = beta * logratios_chunk
+        chosen_rewards_sum = (rewards * preference_labels_chunk.unsqueeze(1)).sum()
+        rejected_rewards_sum = (rewards * (~preference_labels_chunk).unsqueeze(1)).sum()
 
-    @staticmethod
+        return losses.sum() / (full_target.shape[0]), chosen_rewards_sum, rejected_rewards_sum
+
+    @classmethod
     def forward(
+        cls,
         ctx,
         _input,
         weight,
@@ -81,15 +85,38 @@ class LigerFusedLinearKTOFunction(LigerFusedLinearUnpairedPreferenceBase):
         beta=0.1,
         compiled=True,
         use_ref_model=True,
+        average_log_prob=False,
+        chunk_size=1,
     ):
-        return LigerFusedLinearUnpairedPreferenceBase.forward(
+        """
+        Fused linear layer with KTO loss.
+        Args:
+            _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
+            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            target (torch.LongTensor): Target tensor. Shape: (batch_size * seq_len,)
+            preference_labels (torch.Tensor): Preference labels tensor. Shape: (batch_size,)
+            bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ref_input (torch.Tensor, optional): Reference model input tensor. Shape: (batch_size * seq_len, hidden_size)
+            ref_weight (torch.Tensor, optional): Reference model weight tensor. Shape: (vocab_size, hidden_size)
+            ref_bias (torch.Tensor, optional): Reference model bias tensor. Shape: (vocab_size,)
+            kl (torch.Tensor, optional): KL divergence tensor. Shape: (batch_size,)
+            ignore_index (int): Index to ignore in loss computation
+            beta (float): Temperature parameter for the KTO loss
+            compiled (bool): Whether to use torch compile
+            use_ref_model (bool): Whether to use a reference model
+            average_log_prob (bool): Whether to average the log probability per non-masked token
+            chunk_size (int): Size of chunks for processing
+        Returns:
+            torch.Tensor: Computed loss
+        """
+        return super().forward(
+            cls=cls,
             ctx=ctx,
             _input=_input,
             weight=weight,
             target=target,
             preference_labels=preference_labels,
             bias=bias,
-            loss_fn=LigerFusedLinearKTOFunction.preference_loss_fn,
             ignore_index=ignore_index,
             beta=beta,
             compiled=compiled,
@@ -97,7 +124,9 @@ class LigerFusedLinearKTOFunction(LigerFusedLinearUnpairedPreferenceBase):
             ref_input=ref_input,
             ref_weight=ref_weight,
             ref_bias=ref_bias,
+            average_log_prob=average_log_prob,
             kl=kl,
+            chunk_size=chunk_size,
         )
 
     @staticmethod
@@ -115,6 +144,7 @@ class LigerFusedLinearKTOFunction(LigerFusedLinearUnpairedPreferenceBase):
             None,
             None,
             None,
+            None,
         )
 
 
@@ -129,6 +159,8 @@ class LigerFusedLinearKTOLoss(torch.nn.Module):
         beta: float = 0.1,
         compiled: bool = True,
         use_ref_model: bool = False,
+        average_log_prob: bool = False,
+        chunk_size: int = 1,
     ):
         """
         Args:
@@ -136,12 +168,16 @@ class LigerFusedLinearKTOLoss(torch.nn.Module):
             beta (float): Temperature parameter for the KTO loss
             compiled (bool): Whether to use compiled operations
             use_ref_model (bool): Whether to use a reference model for the DPO loss.
+            average_log_prob (bool): Whether to average the log probability per non-masked token
+            chunk_size (int): Size of chunks for processing
         """
         super().__init__()
         self.ignore_index = ignore_index
         self.beta = beta
         self.compiled = compiled
         self.use_ref_model = use_ref_model
+        self.average_log_prob = average_log_prob
+        self.chunk_size = chunk_size
 
     def forward(
         self,
@@ -169,4 +205,6 @@ class LigerFusedLinearKTOLoss(torch.nn.Module):
             self.beta,
             self.compiled,
             self.use_ref_model,
+            self.average_log_prob,
+            self.chunk_size,
         )

liger_kernel/chunked_loss/orpo_loss.py

@@ -42,8 +42,9 @@ class LigerFusedLinearORPOFunction(LigerFusedLinearPreferenceBase):
 
         return loss, chosen_rewards, rejected_rewards, log_odds_ratio, log_odds_chosen
 
-    @staticmethod
+    @classmethod
     def forward(
+        cls,
         ctx,
         _input,
         weight,
@@ -54,25 +55,43 @@ class LigerFusedLinearORPOFunction(LigerFusedLinearPreferenceBase):
         compute_nll_loss=True,
         nll_target=None,
         compiled=True,
+        chunk_size=1,
     ):
-        return LigerFusedLinearPreferenceBase.forward(
+        """
+        Fused linear layer with ORPO loss.
+        Args:
+            _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
+            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            target (torch.LongTensor): Target tensor. Shape: (batch_size * seq_len,)
+            bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ignore_index (int): Index to ignore in loss computation
+            beta (float): Weight for the odds ratio loss
+            compute_nll_loss (bool): Whether to compute the NLL loss
+            nll_target (torch.LongTensor, optional): Target tensor for NLL loss. Shape: (batch_size * seq_len,)
+            compiled (bool): Whether to use torch compile
+            chunk_size (int): Size of chunks for processing
+        Returns:
+            torch.Tensor: Computed loss
+        """
+        return super().forward(
+            cls=cls,
             ctx=ctx,
             _input=_input,
             weight=weight,
             target=target,
             bias=bias,
-            loss_fn=LigerFusedLinearORPOFunction.preference_loss_fn,
             ignore_index=ignore_index,
             beta=beta,
             compute_nll_loss=compute_nll_loss,
             nll_target=nll_target,
             compiled=compiled,
+            chunk_size=chunk_size,
         )
 
     @staticmethod
     def backward(ctx, *grad_output):
         grads = LigerFusedLinearPreferenceBase.backward(ctx, grad_output)[:4]
-        return *grads, None, None, None, None, None
+        return *grads, None, None, None, None, None, None
 
 
 class LigerFusedLinearORPOLoss(torch.nn.Module):
@@ -86,19 +105,31 @@ class LigerFusedLinearORPOLoss(torch.nn.Module):
         beta: float = 0.1,
         compute_nll_loss: bool = True,
         compiled: bool = True,
+        chunk_size: int = 1,
     ):
         """
         Args:
             ignore_index (int): Index to ignore in the loss.
             beta (float): Weight for the odds ratio loss.
+            compute_nll_loss (bool): Whether to compute the NLL loss.
+            compiled (bool): Whether to use the torch compiled kernel.
+            chunk_size (int): Size of chunks for processing.
         """
         super().__init__()
         self.ignore_index = ignore_index
         self.beta = beta
         self.compute_nll_loss = compute_nll_loss
         self.compiled = compiled
+        self.chunk_size = chunk_size
 
-    def forward(self, lin_weight, _input, target, bias=None, nll_target=None):
+    def forward(
+        self,
+        lin_weight,
+        _input,
+        target,
+        bias=None,
+        nll_target=None,
+    ):
         return LigerFusedLinearORPOFunction.apply(
             _input,
             lin_weight,
@@ -109,4 +140,5 @@ class LigerFusedLinearORPOLoss(torch.nn.Module):
             self.compute_nll_loss,
             nll_target,
             self.compiled,
+            self.chunk_size,
         )

liger_kernel/chunked_loss/simpo_loss.py

@@ -47,8 +47,9 @@ class LigerFusedLinearSimPOFunction(LigerFusedLinearPreferenceBase):
 
         return loss, chosen_rewards, rejected_rewards
 
-    @staticmethod
+    @classmethod
     def forward(
+        cls,
         ctx,
         _input,
         weight,
@@ -61,27 +62,47 @@ class LigerFusedLinearSimPOFunction(LigerFusedLinearPreferenceBase):
         compute_nll_loss=False,
         compiled=True,
         gamma=0.5,
+        chunk_size=1,
     ):
-        return LigerFusedLinearPreferenceBase.forward(
-            ctx,
-            _input,
-            weight,
-            target,
-            bias,
-            loss_fn=LigerFusedLinearSimPOFunction.preference_loss_fn,
-            compute_nll_loss=compute_nll_loss,
+        """
+        Fused linear layer with SimPO loss.
+        Args:
+            _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
+            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            target (torch.LongTensor): Target tensor. Shape: (batch_size * seq_len,)
+            bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ignore_index (int): Index to ignore in loss computation
+            beta (float): Weight for the odds ratio loss
+            alpha (float): Weight for the alpha parameter
+            label_smoothing (float): Label smoothing factor
+            compute_nll_loss (bool): Whether to compute the NLL loss
+            compiled (bool): Whether to use torch compile
+            gamma (float): Weight for the gamma parameter
+            chunk_size (int): Size of chunks for processing
+        Returns:
+            torch.Tensor: Computed loss
+        """
+        return super().forward(
+            cls=cls,
+            ctx=ctx,
+            _input=_input,
+            weight=weight,
+            target=target,
+            bias=bias,
             ignore_index=ignore_index,
             alpha=alpha,
             beta=beta,
             label_smoothing=label_smoothing,
+            compute_nll_loss=compute_nll_loss,
             compiled=compiled,
             gamma=gamma,
+            chunk_size=chunk_size,
         )
 
     @staticmethod
     def backward(ctx, *grad_output):
         grads = LigerFusedLinearPreferenceBase.backward(ctx, grad_output)[:4]
-        return *grads, None, None, None, None, None, None, None
+        return *grads, None, None, None, None, None, None, None, None
 
 
 class LigerFusedLinearSimPOLoss(torch.nn.Module):
@@ -98,11 +119,18 @@ class LigerFusedLinearSimPOLoss(torch.nn.Module):
         compute_nll_loss: bool = True,
         compiled: bool = True,
         gamma: float = 0.5,
+        chunk_size: int = 1,
     ):
         """
         Args:
             ignore_index (int): Index to ignore in the loss.
             beta (float): Weight for the odds ratio loss.
+            alpha (float): Weight for the alpha parameter.
+            label_smoothing (float): Label smoothing factor.
+            compute_nll_loss (bool): Whether to compute the NLL loss.
+            compiled (bool): Whether to use the torch compiled kernel.
+            gamma (float): Weight for the gamma parameter.
+            chunk_size (int): Size of chunks for processing.
         """
         super().__init__()
         self.ignore_index = ignore_index
@@ -112,8 +140,15 @@ class LigerFusedLinearSimPOLoss(torch.nn.Module):
         self.compute_nll_loss = compute_nll_loss
         self.compiled = compiled
         self.gamma = gamma
+        self.chunk_size = chunk_size
 
-    def forward(self, lin_weight, _input, target, bias=None):
+    def forward(
+        self,
+        lin_weight,
+        _input,
+        target,
+        bias=None,
+    ):
         return LigerFusedLinearSimPOFunction.apply(
             _input,
             lin_weight,
@@ -126,4 +161,5 @@ class LigerFusedLinearSimPOLoss(torch.nn.Module):
             self.compute_nll_loss,
             self.compiled,
             self.gamma,
+            self.chunk_size,
         )

liger_kernel/ops/cross_entropy.py

@@ -285,13 +285,17 @@ def cross_entropy_forward(
 
     target_mask = target != ignore_index
     n_non_ignore = target_mask.sum().item()
+    assert (target * target_mask).max() < _input.shape[-1], (
+        f"Target {target.max()} is out of bounds. Expected < {_input.shape[-1]}"
+    )
+    assert (target * target_mask).min() >= 0, f"Target {target.min()} is out of bounds. Expected >= 0"
     sum_non_ignore_weight = n_non_ignore
     weight_sum = 0.0
     if weight is not None:
         assert weight.shape[0] == V, f"If given, weight has to be a Tensor of size V. Got: {weight.shape}"
-        assert torch.is_floating_point(
-            weight
-        ), f"If given, weight has to be a Tensor of floating point dtype. Got: {weight.dtype}"
+        assert torch.is_floating_point(weight), (
+            f"If given, weight has to be a Tensor of floating point dtype. Got: {weight.dtype}"
+        )
         sum_non_ignore_weight = torch.gather(weight, dim=0, index=target.masked_select(target_mask)).sum().item()
         weight_sum = weight.sum().item()
         # ensure weight is contiguous

liger_kernel/ops/fused_linear_cross_entropy.py

@@ -58,9 +58,9 @@ def fused_linear_cross_entropy_forward(
     ce_weight_sum = 0.0
     if ce_weight is not None:
         assert ce_weight.shape[0] == V, f"If given, weight has to be a Tensor of size V. Got: {ce_weight.shape}"
-        assert torch.is_floating_point(
-            ce_weight
-        ), f"If given, weight has to be a Tensor of floating point dtype. Got: {ce_weight.dtype}"
+        assert torch.is_floating_point(ce_weight), (
+            f"If given, weight has to be a Tensor of floating point dtype. Got: {ce_weight.dtype}"
+        )
         total_sum_non_ignore_ce_weight = (
             torch.gather(ce_weight, dim=0, index=target.masked_select(target_mask)).sum().item()
         )

liger_kernel/ops/fused_linear_jsd.py

@@ -195,9 +195,9 @@ class LigerFusedLinearJSDFunction(torch.autograd.Function):
         """
         has_label = False
         if shift_labels is not None:
-            assert shift_labels.shape == (
-                teacher_input.shape[0],
-            ), f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            assert shift_labels.shape == (teacher_input.shape[0],), (
+                f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            )
             shift_labels = shift_labels.contiguous()
             has_label = True
 

liger_kernel/ops/jsd.py

@@ -157,9 +157,9 @@ class LigerJSDFunction(torch.autograd.Function):
         """
         has_label = False
         if shift_labels is not None:
-            assert shift_labels.shape == (
-                _input.shape[0],
-            ), f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            assert shift_labels.shape == (_input.shape[0],), (
+                f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            )
             shift_labels = shift_labels.contiguous()
             has_label = True
 

liger_kernel/ops/layer_norm.py

@@ -57,13 +57,14 @@ def _layer_norm_forward_kernel(
     B_row = tl.load(B_ptr + col_offsets, mask=mask, other=0)
 
     mean = tl.sum(X_row, axis=0) / n_cols
-    var = tl.sum((X_row - mean) * (X_row - mean), axis=0) / n_cols
+    Xmm = tl.where(mask, X_row - mean, 0)
+    var = tl.sum(Xmm * Xmm, axis=0) / n_cols
     rstd = rsqrt(var + eps)
 
     tl.store(Mean_ptr, mean)
     tl.store(RSTD_ptr, rstd)
 
-    Y_row = (X_row - mean) * rstd * W_row + B_row
+    Y_row = Xmm * rstd * W_row + B_row
 
     tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
 
@@ -147,9 +148,11 @@ def layer_norm_forward(X, W, B, eps):
     Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
     Mean = torch.empty(n_rows, dtype=X.dtype, device=X.device)
     RSTD = torch.empty(n_rows, dtype=X.dtype, device=X.device)
-    assert (
-        X.shape[1] == W.shape[0]
-    ), f"Incompatible hidden size dimension between input tensor with shape[1] = {X.shape[1]} and weight tensor with shape[0] = {W.shape[0]}"
+    if X.shape[1] != W.shape[0]:
+        raise ValueError(
+            f"Incompatible dimensions: input feature size (X.shape[1]={X.shape[1]}) "
+            f"must match weight size (W.shape[0]={W.shape[0]})"
+        )
 
     _layer_norm_forward_kernel[(n_rows,)](
         Y,
@@ -190,11 +193,21 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
 
     BLOCK_SIZE, num_warps = calculate_settings(n_cols)
     if n_cols > BLOCK_SIZE:
-        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+        raise RuntimeError(
+            f"Feature dimension {n_cols} exceeds maximum supported size of {BLOCK_SIZE}. Consider using a smaller feature dimension."
+        )
 
     rows_per_program = math.ceil(n_rows / sm_count)
     grid = (sm_count,)
-    triton_dtype = tl.float32 if X.dtype == torch.float32 else tl.bfloat16
+    triton_dtype = (
+        tl.float32
+        if X.dtype == torch.float32
+        else tl.bfloat16
+        if X.dtype == torch.bfloat16
+        else tl.float16
+        if X.dtype == torch.float16
+        else tl.float32  # fallback to float32 for other types
+    )
     _layer_norm_backward_kernel[grid](
         X,
         W,

liger_kernel/ops/tvd.py

@@ -0,0 +1,207 @@
+from typing import Literal
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.utils import ensure_contiguous
+
+MAX_FUSED_SIZE = 65536 // 4
+
+REDUCTION_LITERAL = Literal["none", "sum", "mean", "batchmean"]
+
+_REDUCTION_MODE_NONE = tl.constexpr(0)
+_REDUCTION_MODE_SUM = tl.constexpr(1)
+_REDUCTION_MODE_MEAN = tl.constexpr(2)
+_REDUCTION_MODE_BATCHMEAN = tl.constexpr(3)
+
+_str_to_reduction_mode = {
+    "none": _REDUCTION_MODE_NONE.value,
+    "sum": _REDUCTION_MODE_SUM.value,
+    "mean": _REDUCTION_MODE_MEAN.value,
+    "batchmean": _REDUCTION_MODE_BATCHMEAN.value,
+}
+
+
+def get_num_warps(BLOCK_SIZE):
+    num_warps = 4
+    if BLOCK_SIZE >= 32768:
+        num_warps = 32
+    elif BLOCK_SIZE >= 8192:
+        num_warps = 16
+    elif BLOCK_SIZE >= 2048:
+        num_warps = 8
+
+    return num_warps
+
+
+@triton.jit
+def _tv_distance_kernel(
+    p_ptr,
+    p_stride,
+    q_ptr,
+    q_stride,
+    loss_ptr,
+    loss_stride,
+    grads_ptr,
+    grads_stride,
+    label_ptr,
+    ignore_index: tl.constexpr,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+    HAS_LABEL: tl.constexpr,
+    reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    p_ptr += pid * p_stride
+    q_ptr += pid * q_stride
+    loss_ptr += pid * loss_stride
+    grads_ptr += pid * grads_stride
+    label_ptr += pid
+
+    base_offsets = tl.arange(0, BLOCK_SIZE)
+
+    if HAS_LABEL:
+        label = tl.load(label_ptr)
+        if label == ignore_index:
+            for i in range(0, n_cols, BLOCK_SIZE):
+                offsets = i + base_offsets
+                mask = offsets < n_cols
+                tl.store(grads_ptr + offsets, 0.0, mask=mask)
+                if reduction == _REDUCTION_MODE_NONE:
+                    tl.store(loss_ptr + offsets, 0.0, mask=mask)
+            return
+
+    loss_sum = 0.0
+    for i in range(0, n_cols, BLOCK_SIZE):
+        offsets = i + base_offsets
+        mask = offsets < n_cols
+
+        p = tl.load(p_ptr + offsets, mask=mask, other=0.0)
+        q = tl.load(q_ptr + offsets, mask=mask, other=0.0)
+
+        # TVD(P || Q) = 0.5 * |P - Q|
+        tv_loss = 0.5 * tl.abs(p - q)
+
+        grad_res = tl.where(p > q, 0.5, -0.5)
+
+        tl.store(grads_ptr + offsets, grad_res, mask=mask)
+
+        if reduction == _REDUCTION_MODE_NONE:
+            tl.store(loss_ptr + offsets, tv_loss, mask=mask)
+        else:
+            loss_sum += tl.sum(tv_loss, axis=0)
+
+    if reduction != _REDUCTION_MODE_NONE:
+        tl.store(loss_ptr, loss_sum)
+
+
+def tv_distance_forward_triton(p, q, shift_labels, reduction, ignore_index, has_label):
+    BT, V = p.shape
+
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+    num_warps = get_num_warps(BLOCK_SIZE)
+
+    grid = (BT,)
+
+    reduction = _str_to_reduction_mode[reduction]
+
+    out_size = (BT, V) if reduction == _REDUCTION_MODE_NONE.value else (BT,)
+    output_tensor = torch.zeros(out_size, device=p.device, dtype=torch.float32)
+    grads = torch.empty_like(p)
+
+    n_non_ignore = (shift_labels != ignore_index).sum().item() if has_label else BT
+
+    _tv_distance_kernel[grid](
+        p,
+        p.stride(0),
+        q,
+        q.stride(0),
+        output_tensor,
+        output_tensor.stride(0),
+        grads,
+        grads.stride(0),
+        shift_labels if has_label else torch.empty(1, device=p.device),
+        ignore_index,
+        V,
+        BLOCK_SIZE=BLOCK_SIZE,
+        HAS_LABEL=has_label,
+        num_warps=num_warps,
+        reduction=reduction,
+    )
+
+    if reduction == _REDUCTION_MODE_BATCHMEAN.value:
+        return output_tensor.sum() / n_non_ignore, grads / n_non_ignore
+    elif reduction == _REDUCTION_MODE_SUM.value:
+        return output_tensor.sum(dim=0), grads
+    elif reduction == _REDUCTION_MODE_MEAN.value:
+        return output_tensor.sum() / (n_non_ignore * V), grads / (n_non_ignore * V)
+    else:
+        return output_tensor, grads
+
+
+def tvd_backward_triton(grad_output, grads):
+    # If cross entropy is the last layer, grad_output is 1.0. Skip the mul then.
+    if torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)):
+        return grads
+
+    return grads * grad_output
+
+
+class LigerTVDLossFunction(torch.autograd.Function):
+    """
+    Class implementing the forward and backward pass for the Total Variation Distance Loss using Triton.
+    """
+
+    @staticmethod
+    @ensure_contiguous
+    def forward(
+        ctx,
+        p: torch.Tensor,
+        q: torch.Tensor,
+        shift_labels: Optional[torch.Tensor] = None,
+        reduction: REDUCTION_LITERAL = "batchmean",
+        ignore_index: int = -100,
+    ) -> torch.Tensor:
+        """A forward pass for the Total Variation Distance Loss.
+
+        Args:
+            ctx: Torch autograd context
+            p (torch.Tensor): A tensor of shape (BT, V) containing the first distribution.
+            q (torch.Tensor): A tensor of shape (BT, V) containing the second distribution.
+            shift_labels (Optional[torch.Tensor]): A tensor of shape (BT,) containing the labels.
+            reduction (REDUCTION_LITERAL, optional): The reduction method to be applied. Defaults to "batchmean".
+            ignore_index (int, optional): The index to ignore during loss calculation. Defaults to -100.
+
+        Returns:
+            torch.Tensor: The computed Total Variation Distance Loss.
+        """
+        has_label = False
+        if shift_labels is not None:
+            assert shift_labels.shape == (p.shape[0],), (
+                f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            )
+            shift_labels = shift_labels.contiguous()
+            has_label = True
+
+        loss, grads = tv_distance_forward_triton(p, q, shift_labels, reduction, ignore_index, has_label)
+        ctx.save_for_backward(grads)
+        return loss
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+        """A backward pass for the Total Variation Distance Loss.
+
+        Args:
+            ctx: Torch autograd context
+            grad_output (torch.Tensor): The gradient of the loss with respect to the output.
+
+        Returns:
+            tuple[torch.Tensor, None, None, None, None]: The gradient of the loss with respect to the inputs.
+        """
+        (grads,) = ctx.saved_tensors
+        grads = tvd_backward_triton(grad_output, grads)
+
+        return grads, None, None, None, None