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

liger_kernel/chunked_loss/grpo_loss.py

@@ -1,66 +1,92 @@
 import torch
 
-from liger_kernel.chunked_loss.fused_linear_rlhf import LigerFusedLinearRLHFBase
+from liger_kernel.chunked_loss.fused_linear_ppo import LigerFusedLinearPPOBase
 
 
-class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
+def k3_loss_fn(log_p, log_q):
+    # computes k3 estimate of KL[q, p]
+    # ref: http://joschu.net/blog/kl-approx.html
+    return torch.exp(log_p - log_q) - (log_p - log_q) - 1.0
+
+
+def clip_coef_fn(coef, epsilon_low, epsilon_high):
+    return torch.clamp(coef, 1 - epsilon_low, 1 + epsilon_high)
+
+
+class LigerFusedLinearGRPOFunction(LigerFusedLinearPPOBase):
     @staticmethod
-    def rlhf_loss_fn(
+    def ppo_loss_fn(
         log_probs,
+        selected_token_ids,
         attention_mask,
-        rewards,
-        ref_log_probs=None,
-        beta=0.1,
+        advantages,
+        full_attention_mask,
+        ref_per_token_logps=None,  # shape: [chunk_size, seq_len]
+        old_per_token_logps=None,
+        ref_log_probs=None,  # used when ref_per_token_logps is None (shape: [chunk_size, seq_len, vocab_size])
+        epsilon_low=0.2,
+        epsilon_high=0.2,
+        beta=0.04,
+        loss_type="bnpo",  # ["grpo", "bnpo", "dr_grpo"]
+        max_completion_length=None,  # Required for dr_grpo
         **kwargs,
     ):
         """GRPO Loss Function matching GRPOTrainer implementation."""
-        # Get chosen token probabilities
-        chosen_tokens = log_probs.argmax(dim=-1)  # (batch_size, seq_len)
-        chosen_token_logprobs = log_probs.gather(dim=-1, index=chosen_tokens.unsqueeze(-1)).squeeze(
+        per_token_logps = log_probs.gather(dim=-1, index=selected_token_ids.unsqueeze(-1)).squeeze(
             -1
         )  # (batch_size, seq_len)
 
         # Get reference model probabilities
-        if ref_log_probs is not None:
-            with torch.no_grad():
-                ref_token_logprobs = ref_log_probs.gather(dim=-1, index=chosen_tokens.unsqueeze(-1)).squeeze(-1)
-        else:
-            ref_token_logprobs = chosen_token_logprobs.detach()
-
-        # Compute advantages per batch entry in a grouped fashion
-        mean_grouped_rewards = rewards.mean()  # [batch_size,]
-        std_grouped_rewards = rewards.std()  # [batch_size,]
-
-        # Calculate advantages using the same epsilon as in GRPOTrainer
-        eps = 1e-4
-        advantages = (rewards - mean_grouped_rewards) / (std_grouped_rewards + eps)
+        if ref_per_token_logps is None:
+            if ref_log_probs is not None:
+                with torch.no_grad():
+                    ref_per_token_logps = ref_log_probs.gather(dim=-1, index=selected_token_ids.unsqueeze(-1)).squeeze(
+                        -1
+                    )
+            else:
+                ref_per_token_logps = per_token_logps.detach()
 
         # Compute policy gradient loss with importance sampling ratio
-        ratio = torch.exp(chosen_token_logprobs - chosen_token_logprobs.detach())
-        policy_loss = -ratio * advantages.unsqueeze(1)
-
-        # Compute KL penalty
-        kl_div = (
-            torch.exp(ref_token_logprobs - chosen_token_logprobs) - (ref_token_logprobs - chosen_token_logprobs) - 1.0
-        )
+        old_per_token_logps = old_per_token_logps if old_per_token_logps is not None else per_token_logps.detach()
+        coef_1 = torch.exp(per_token_logps - old_per_token_logps)
+        coef_2 = clip_coef_fn(coef_1, epsilon_low, epsilon_high)
+        per_token_loss1 = coef_1 * advantages.unsqueeze(1)
+        per_token_loss2 = coef_2 * advantages.unsqueeze(1)
+        per_token_loss = -torch.min(per_token_loss1, per_token_loss2)
+        if beta != 0.0:
+            # Compute KL penalty (approximates KL[per_token_logps, ref_per_token_logps])
+            kl_div = k3_loss_fn(ref_per_token_logps, per_token_logps)
+            # Combine losses
+            per_token_loss = per_token_loss + beta * kl_div
 
-        # Combine losses
-        per_token_loss = policy_loss + beta * kl_div
-
-        # Apply masking and normalize
-        masked_loss = per_token_loss * attention_mask
-        seq_lengths = attention_mask.sum()
-        seq_lengths = torch.clamp(seq_lengths, min=1.0)
-        loss = masked_loss.sum() / seq_lengths
+        # Note: We normalize by the number of tokens in the batch (using full_attention_mask),
+        # which is consistent with the DAPO loss implementation (https://arxiv.org/html/2503.14476v1)
+        # and TRL GRPO implementation
+        # (https://github.com/huggingface/trl/blob/e751a16df56e70190fb94bed4a2035eec3303777/trl/trainer/grpo_trainer.py#L966)
+        if loss_type == "grpo":
+            # Average per-sequence loss
+            loss = (
+                (per_token_loss * attention_mask).sum(-1) / torch.clamp(attention_mask.sum(-1), min=1.0)
+            ).sum() / full_attention_mask.shape[0]
+        elif loss_type == "bnpo":
+            # Batch Normalized Per-token loss (original implementation)
+            loss = (per_token_loss * attention_mask).sum() / torch.clamp(full_attention_mask.sum(), min=1.0)
+        elif loss_type == "dr_grpo":
+            # Dimension-Reduced GRPO (normalize by batch_size * max_completion_length)
+            if max_completion_length is None:
+                raise ValueError("max_completion_length must be provided for loss_type 'dr_grpo'")
+            loss = (per_token_loss * attention_mask).sum() / (full_attention_mask.shape[0] * max_completion_length)
+        else:
+            raise ValueError(f"Unknown loss type: {loss_type}")
 
         # Calculate metrics
-        metrics = (
-            chosen_token_logprobs.mean(),  # mean log prob
-            chosen_token_logprobs.std(),  # std log prob
-            log_probs.mean(),  # mean all log probs
-            ((kl_div * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)).mean(),  # mean KL div
+        metrics = []
+        if beta != 0.0:
+            metrics.append(((kl_div * attention_mask).sum() / torch.clamp(full_attention_mask.sum(), min=1.0)))
+        is_clipped = ((coef_1 < 1 - epsilon_low) & (advantages.unsqueeze(1) < 0)) | (
+            (coef_1 > 1 + epsilon_high) & (advantages.unsqueeze(1) > 0)
         )
-
+        metrics.append((is_clipped * attention_mask).sum() / torch.clamp(full_attention_mask.sum(), min=1.0))
         return loss, metrics
 
     @classmethod
@@ -69,16 +95,23 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
         ctx,
         _input,
         weight,
+        selected_token_ids,
         attention_mask,
-        rewards,
+        advantages,
         bias=None,
+        ref_per_token_logps=None,
+        old_per_token_logps=None,
         ref_input=None,
         ref_weight=None,
         ref_bias=None,
-        beta=0.1,
+        beta=0.04,
+        epsilon_low=0.2,
+        epsilon_high=0.2,
+        loss_type="bnpo",
+        max_completion_length=None,
+        temperature=1.0,
         compiled=True,
         use_ref_model=True,
-        num_generations=1,
         chunk_size=1,
     ):
         """
@@ -86,16 +119,20 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
         Args:
             _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
             weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            selected_token_ids (torch.Tensor): Selected token ids tensor. Shape: (batch_size, seq_len)
             attention_mask (torch.Tensor): Attention mask tensor. Shape: (batch_size, seq_len)
-            rewards (torch.Tensor): Rewards tensor. Shape: (batch_size,)
+            advantages (torch.Tensor): Advantages tensor. Shape: (batch_size,)
             bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ref_per_token_logps:  Reference model log probs per token tensor. Shape:(batch_size, seq_len)
             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,)
             beta (float): Weight for the KL penalty
+            loss_type (str): Type of loss calculation ("grpo", "bnpo", "dr_grpo"). Defaults to "bnpo".
+            max_completion_length (int, optional): Maximum completion length, required for "dr_grpo". Defaults to None.
+            temperature (float): Temperature for the logits
             compiled (bool): Whether to use torch compile
             use_ref_model (bool): Whether to use a reference model
-            num_generations (int): Number of generations per prompt
             chunk_size (int): Size of chunks for processing.
         Returns:
             torch.Tensor: Computed loss
@@ -105,16 +142,23 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
             ctx=ctx,
             _input=_input,
             weight=weight,
+            selected_token_ids=selected_token_ids,
             attention_mask=attention_mask,
-            rewards=rewards,
+            advantages=advantages,
             bias=bias,
+            ref_per_token_logps=ref_per_token_logps,
+            old_per_token_logps=old_per_token_logps,
             ref_input=ref_input,
             ref_weight=ref_weight,
             ref_bias=ref_bias,
             beta=beta,
+            epsilon_low=epsilon_low,
+            epsilon_high=epsilon_high,
+            loss_type=loss_type,
+            max_completion_length=max_completion_length,
+            temperature=temperature,
             compiled=compiled,
             use_ref_model=use_ref_model,
-            num_generations=num_generations,
             chunk_size=chunk_size,
         )
 
@@ -126,16 +170,24 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
             grad_output: Gradient of the loss (scalar)
             grad_metrics: Gradients of the metrics (not used in backward computation)
         """
-        grads = LigerFusedLinearRLHFBase.backward(ctx, grad_output)
+        grads = LigerFusedLinearPPOBase.backward(ctx, grad_output)
         return (
-            *grads[:5],  # grad_input, grad_weight, grad_attention_mask, grad_rewards, grad_bias
+            *grads[
+                :6
+            ],  # grad_input, grad_weight, grad_selected_token_ids, grad_attention_mask, grad_advantages, grad_bias
+            None,  # grad_ref_per_token_logps
+            None,  # grad_old_per_token_logps
             None,  # grad_ref_input
             None,  # grad_ref_weight
             None,  # grad_ref_bias
             None,  # grad_beta
+            None,  # grad_epsilon_low
+            None,  # grad_epsilon_high
+            None,  # grad_loss_type (string, not differentiable)
+            None,  # grad_max_completion_length (int, not differentiable)
+            None,  # grad_temperature
             None,  # grad_compiled
             None,  # grad_use_ref_model
-            None,  # grad_num_generations
             None,  # grad_chunk_size
         )
 
@@ -145,34 +197,49 @@ class LigerFusedLinearGRPOLoss(torch.nn.Module):
 
     def __init__(
         self,
-        beta: float = 0.1,
+        beta: float = 0.04,
         compiled: bool = True,
         use_ref_model: bool = True,
-        num_generations: int = 1,
         chunk_size: int = 1,
+        epsilon_low: float = 0.2,
+        epsilon_high: float = 0.2,
+        loss_type: str = "bnpo",
+        max_completion_length: int | None = None,
+        temperature: float = 1.0,
     ):
         """
         Args:
             beta (float): Weight for the KL penalty.
             compiled (bool): Whether to use torch compile.
             use_ref_model (bool): Whether to use a reference model.
-            num_generations (int): Number of generations per prompt.
             chunk_size (int): Size of chunks for processing.
+            epsilon_low (float): Lower bound for the importance sampling ratio.
+            epsilon_high (float): Upper bound for the importance sampling ratio.
+            loss_type (str): Type of loss calculation ("grpo", "bnpo", "dr_grpo"). Defaults to "bnpo".
+            max_completion_length (int, optional): Maximum completion length, required for "dr_grpo". Defaults to None.
+            temperature (float): Temperature for the logits.
         """
         super().__init__()
         self.beta = beta
         self.compiled = compiled
         self.use_ref_model = use_ref_model
-        self.num_generations = num_generations
         self.chunk_size = chunk_size
+        self.epsilon_low = epsilon_low
+        self.epsilon_high = epsilon_high
+        self.loss_type = loss_type
+        self.max_completion_length = max_completion_length
+        self.temperature = temperature
 
     def forward(
         self,
         _input,
         lin_weight,
+        selected_token_ids,
         attention_mask,
-        rewards,
+        advantages,
         bias=None,
+        ref_per_token_logps=None,
+        old_per_token_logps=None,
         ref_input=None,
         ref_weight=None,
         ref_bias=None,
@@ -180,15 +247,22 @@ class LigerFusedLinearGRPOLoss(torch.nn.Module):
         return LigerFusedLinearGRPOFunction.apply(
             _input,
             lin_weight,
+            selected_token_ids,
             attention_mask,
-            rewards,
+            advantages,
             bias,
+            ref_per_token_logps,
+            old_per_token_logps,
             ref_input,
             ref_weight,
             ref_bias,
             self.beta,
+            self.epsilon_low,
+            self.epsilon_high,
+            self.loss_type,
+            self.max_completion_length,
+            self.temperature,
             self.compiled,
             self.use_ref_model,
-            self.num_generations,
             self.chunk_size,
         )

liger_kernel/chunked_loss/jsd_loss.py

@@ -19,15 +19,20 @@ class LigerFusedLinearJSDFunction(LigerFusedLinearDistillationBase):
         student_log_probs = F.log_softmax(student_logits, dim=-1)
         teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)
 
-        # Compute probabilities (only required for mean calculation)
-        mean_probs = beta * student_log_probs.exp() + (1 - beta) * teacher_log_probs.exp()
-        log_mean_probs = mean_probs.log()
+        if beta == 0:
+            jsd_loss = F.kl_div(student_log_probs, teacher_log_probs, reduction="sum", log_target=True)
+        elif beta == 1:
+            jsd_loss = F.kl_div(teacher_log_probs, student_log_probs, reduction="sum", log_target=True)
+        else:
+            # Compute probabilities (only required for mean calculation)
+            mean_probs = (1 - beta) * student_log_probs.exp() + beta * teacher_log_probs.exp()
+            log_mean_probs = mean_probs.log()
 
-        student_kl = F.kl_div(log_mean_probs, student_log_probs, reduction="sum", log_target=True)
-        teacher_kl = F.kl_div(log_mean_probs, teacher_log_probs, reduction="sum", log_target=True)
+            student_kl = F.kl_div(log_mean_probs, student_log_probs, reduction="sum", log_target=True)
+            teacher_kl = F.kl_div(log_mean_probs, teacher_log_probs, reduction="sum", log_target=True)
 
-        # JSD is the weighted average of the KL divergences
-        jsd_loss = beta * teacher_kl + (1 - beta) * student_kl
+            # JSD is the weighted average of the KL divergences
+            jsd_loss = beta * teacher_kl + (1 - beta) * student_kl
         return jsd_loss
 
     @classmethod

liger_kernel/ops/cross_entropy.py

@@ -9,6 +9,7 @@ import triton.language as tl
 from liger_kernel.ops.utils import compare_version
 from liger_kernel.ops.utils import element_mul_kernel
 from liger_kernel.ops.utils import is_hip
+from liger_kernel.utils import infer_device
 
 if compare_version("triton", operator.ge, "3.0.0"):
     try:
@@ -59,7 +60,7 @@ def liger_cross_entropy_kernel(
     z_loss_ptr: Pointer to tensor to store the z loss. No operation if RETURN_Z_LOSS is 0.
     loss_stride (int): The stride of the loss tensor.
     n_cols (int): The number of columns in the input tensor.
-    n_non_ignore (flaot): The number of non-ignored elements in the batch.
+    n_non_ignore (float): The number of non-ignored elements in the batch.
     sum_non_ignore_weight (float): The sum of non-ignored target's weights in the batch.
     weight_sum (float): The sum of weight tensor.
     ignore_index (int): The index to ignore in the target.
@@ -258,7 +259,7 @@ def liger_cross_entropy_kernel(
 # The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
 # However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
 # The optimal maximum block size depends on your hardware, your kernel, and your dtype
-MAX_FUSED_SIZE = 65536 // 2  # the best size we found by manually tuning
+MAX_FUSED_SIZE = 4096 if infer_device() == "xpu" else 65536 // 2  # the best size we found by manually tuning
 
 
 def cross_entropy_forward(

liger_kernel/ops/dyt.py

@@ -0,0 +1,225 @@
+import operator
+
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.utils import calculate_settings
+from liger_kernel.ops.utils import compare_version
+from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.ops.utils import infer_device
+
+if compare_version("triton", operator.ge, "3.0.0"):
+    try:
+        # typical import path with dispatch available
+        from triton.language.extra.libdevice import tanh
+    except ModuleNotFoundError:
+        # for working with NGC containers
+        from triton.language.extra.cuda.libdevice import tanh
+else:
+    from triton.language.math import tanh
+
+
+@triton.jit
+def _dyt_fwd_kernel(
+    x_ptr,
+    x_row_stride,
+    alpha_ptr,
+    gamma_ptr,
+    beta_ptr,
+    y_ptr,
+    y_row_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Reference:
+    https://arxiv.org/abs/2503.10622
+
+    Shapes:
+        - x: (BT, C)
+        - alpha: (1)
+        - gamma: (C)
+        - beta: (C)
+    """
+    row_idx = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_cols
+
+    x_ptr += row_idx * x_row_stride
+    y_ptr += row_idx * y_row_stride
+
+    alpha = tl.load(alpha_ptr)
+    gamma = tl.load(gamma_ptr + offsets, mask=mask)
+    beta = tl.load(beta_ptr + offsets, mask=mask)
+    x = tl.load(x_ptr + offsets, mask=mask)
+    y = gamma * tanh((alpha * x).cast(tl.float32)) + beta
+    tl.store(y_ptr + offsets, y, mask=mask)
+
+
+@triton.jit
+def _dyt_bwd_kernel(
+    x_ptr,
+    x_row_stride,
+    dy_ptr,
+    dy_row_stride,
+    dx_ptr,
+    dx_row_stride,
+    alpha_ptr,
+    dalpha_ptr,
+    gamma_ptr,
+    dgamma_ptr,
+    dgamma_row_stride,
+    n_cols,
+    n_rows,
+    ROWS_PER_PROGRAM: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Reference:
+    https://arxiv.org/abs/2503.10622
+
+    Shapes:
+        - x: (BT, C)
+        - alpha: (1)
+        - gamma: (C)
+        - dx: (BT, C)
+        - dy: (BT, C)
+        - dgamma: (sm_count, C)
+        - dalpha: (sm_count,)
+    """
+    # d(gamma * tanh(alpha * x) + beta) / dx
+    # = gamma * (1 - tanh^2(alpha * x)) * alpha
+    # d(gamma * tanh(alpha * x) + beta) / dalpha
+    # = gamma * (1 - tanh^2(alpha * x)) * x
+    # d(gamma * tanh(alpha * x) + beta) / dgamma
+    # = tanh(alpha * x)
+    # d(gamma * tanh(alpha * x)) / dbeta = 1
+    pid = tl.program_id(0)
+
+    row_start = pid * ROWS_PER_PROGRAM
+    row_end = min((pid + 1) * ROWS_PER_PROGRAM, n_rows)
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_cols
+
+    dalpha = 0.0
+    dgamma = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+
+    x_ptr += row_start * x_row_stride
+    dx_ptr += row_start * dx_row_stride
+    dy_ptr += row_start * dy_row_stride
+    alpha = tl.load(alpha_ptr)
+    gamma = tl.load(gamma_ptr + offsets, mask=mask, other=0.0)
+
+    for _ in tl.range(row_start, row_end):
+        dy = tl.load(dy_ptr + offsets, mask=mask, other=0.0)
+        x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
+        tanh_ax = tanh((alpha * x).cast(tl.float32))
+        sech2_ax = 1 - tanh_ax * tanh_ax
+
+        dx = dy * gamma * sech2_ax * alpha
+        dalpha += tl.sum(dy * gamma * sech2_ax * x)
+        dgamma += dy * tanh_ax
+        tl.store(dx_ptr + offsets, dx, mask=mask)
+
+        dy_ptr += dy_row_stride
+        x_ptr += x_row_stride
+        dx_ptr += dx_row_stride
+
+    tl.store(dgamma_ptr + pid * dgamma_row_stride + offsets, dgamma, mask=mask)
+    tl.store(dalpha_ptr + pid, dalpha)
+
+    pass
+
+
+def liger_dyt_fwd(x, alpha, gamma, beta):
+    shape = x.shape
+    dim = shape[-1]
+    x = x.view(-1, dim)
+    n_rows, n_cols = x.shape
+    y = torch.empty_like(x)
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    _dyt_fwd_kernel[(n_rows,)](
+        x_ptr=x,
+        alpha_ptr=alpha,
+        gamma_ptr=gamma,
+        beta_ptr=beta,
+        y_ptr=y,
+        x_row_stride=x.stride(0),
+        y_row_stride=y.stride(0),
+        n_cols=n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return y.view(*shape)
+
+
+def liger_dyt_bwd(dy, x, alpha, gamma):
+    shape = dy.shape
+    dtype = x.dtype
+    dim = shape[-1]
+    dy = dy.view(-1, dim)
+    x = x.view(-1, dim)
+    n_rows, n_cols = dy.shape
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    sm_count = 1
+    device = infer_device()
+    if device == "cuda":
+        sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count
+    elif device == "xpu":
+        sm_count = torch.xpu.get_device_properties(x.device).gpu_subslice_count
+    if n_cols > BLOCK_SIZE:
+        raise RuntimeError(
+            f"Feature dimension {dim} exceeds maximum supported size of {BLOCK_SIZE}. Consider using a smaller feature dimension."
+        )
+
+    dx = torch.empty_like(x, dtype=torch.float32)
+    _dalpha = torch.empty((sm_count,), dtype=torch.float32, device=x.device)
+    _dgamma = torch.empty((sm_count, n_cols), dtype=torch.float32, device=x.device)
+
+    grid = (sm_count,)
+    rows_per_program = triton.cdiv(n_rows, sm_count)
+    _dyt_bwd_kernel[grid](
+        x_ptr=x,
+        x_row_stride=x.stride(0),
+        dy_ptr=dy,
+        dy_row_stride=dy.stride(0),
+        dx_ptr=dx,
+        dx_row_stride=dx.stride(0),
+        alpha_ptr=alpha,
+        dalpha_ptr=_dalpha,
+        gamma_ptr=gamma,
+        dgamma_ptr=_dgamma,
+        dgamma_row_stride=_dgamma.stride(0),
+        n_cols=n_cols,
+        n_rows=n_rows,
+        ROWS_PER_PROGRAM=rows_per_program,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    dalpha = _dalpha.sum(dim=0, keepdim=True).to(dtype)
+    dgamma = _dgamma.sum(dim=0).to(dtype)
+    dbeta = dy.sum(dim=0).to(dtype)
+    return dx.view(*shape), dalpha, dgamma, dbeta
+
+
+class LigerDyTFunction(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(ctx, x, alpha, gamma, beta):
+        y = liger_dyt_fwd(x, alpha, gamma, beta)
+        ctx.save_for_backward(x, alpha, gamma)
+        return y
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, grad_output):
+        x, alpha, gamma = ctx.saved_tensors
+        dx, dalpha, dgamma, dbeta = liger_dyt_bwd(
+            grad_output,
+            x,
+            alpha,
+            gamma,
+        )
+
+        return (dx, dalpha, dgamma, dbeta)

liger_kernel/ops/fused_linear_jsd.py

@@ -8,11 +8,12 @@ from liger_kernel.ops.utils import amp_custom_bwd
 from liger_kernel.ops.utils import amp_custom_fwd
 from liger_kernel.ops.utils import element_mul_kernel
 from liger_kernel.ops.utils import is_hip
+from liger_kernel.utils import infer_device
 
 # The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
 # However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
 # The optimal maximum block size depends on your hardware, your kernel, and your dtype
-MAX_FUSED_SIZE = 65536 // 2
+MAX_FUSED_SIZE = 4096 if infer_device() == "xpu" else 65536 // 2
 
 
 def fused_linear_jsd_forward(

liger_kernel/ops/jsd.py

@@ -5,6 +5,7 @@ import triton
 import triton.language as tl
 
 from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.utils import infer_device
 
 
 @triton.jit
@@ -51,29 +52,48 @@ def _jsd_kernel(
         Y = tl.load(Y_ptr + offsets, mask=mask, other=float("-inf")).to(tl.float32)
 
         if beta == 0.0:  # forward KL
-            Y_prob = tl.exp(Y)
+            Y_max = tl.max(Y, axis=0)
+            Y_shifted = Y - Y_max
+            Y_prob = tl.exp(Y_shifted) * tl.exp(Y_max)  # Compensate for the shift
             loss = Y_prob * (Y - X)
             dX = -Y_prob
-        elif beta == 1.0:
-            X_prob = tl.exp(X)
+        elif beta == 1.0:  # reverse KL
+            X_max = tl.max(X, axis=0)
+            X_shifted = X - X_max
+            X_prob = tl.exp(X_shifted) * tl.exp(X_max)  # Compensate for the shift
             loss = X_prob * (X - Y)
             dX = loss + X_prob
         else:
-            Q = tl.exp(X)
-            P = tl.exp(Y)
-            M = beta * P + (1 - beta) * Q
-            log_M = tl.log(M)
+            max_val = tl.maximum(tl.max(X, axis=0), tl.max(Y, axis=0))
+            X_shifted = X - max_val
+            Y_shifted = Y - max_val
 
-            loss = beta * P * Y + (1 - beta) * Q * X - M * log_M
-            dX = (1 - beta) * Q * (X - log_M)
+            # Pre-compute exp(max_val) since it's used twice
+            exp_max = tl.exp(max_val)
+
+            # Compute exp terms with compensation
+            Q = tl.exp(X_shifted) * exp_max  # = exp(X)
+            P = tl.exp(Y_shifted) * exp_max  # = exp(Y)
+
+            # Pre-compute common terms
+            beta_P = beta * P
+            one_minus_beta_Q = (1 - beta) * Q
+            M = beta_P + one_minus_beta_Q
+            log_M = tl.log(M)  # No need to compensate as M is already in original scale
+
+            loss = beta_P * Y + one_minus_beta_Q * X - M * log_M
+            dX = one_minus_beta_Q * (X - log_M)
+
+        # Pre-compute scaling factor
+        scale = 1.0 / n_non_ignore
+        loss = loss * scale
+        dX = dX * scale
 
-        loss = loss / n_non_ignore
-        dX = dX / n_non_ignore
         tl.store(loss_ptr + offsets, loss, mask=mask)
         tl.store(dX_ptr + offsets, dX, mask=mask)
 
 
-MAX_FUSED_SIZE = 65536
+MAX_FUSED_SIZE = 4096 if infer_device() == "xpu" else 65536
 
 
 def jsd_forward(_input, target, shift_labels, beta, ignore_index, has_label):