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

liger_kernel/ops/kl_div.py

@@ -6,6 +6,7 @@ import triton.language as tl
 
 from liger_kernel.ops.utils import ensure_contiguous
 from liger_kernel.ops.utils import is_hip
+from liger_kernel.utils import infer_device
 
 
 def get_num_warps(BLOCK_SIZE):
@@ -115,9 +116,12 @@ def _kldiv_kernel_backward(
 
 def kldiv_forward_triton(y_pred, y_true, log_target, reduction, eps):  # [BT, V]
     BT, V = y_pred.shape
-
-    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
-    num_warps = get_num_warps(BLOCK_SIZE)
+    BLOCK_SIZE = (
+        min(8192, triton.next_power_of_2(V))
+        if infer_device() == "xpu"
+        else min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+    )
+    num_warps = 32 if infer_device() == "xpu" else get_num_warps(BLOCK_SIZE)
 
     grid = (BT,)
     reduction = _str_to_reduction_mode[reduction]
@@ -155,9 +159,12 @@ def kldiv_forward_triton(y_pred, y_true, log_target, reduction, eps):  # [BT, V]
 
 def kldiv_backward_triton(target, grad_output, new_grads, log_target):
     BT, V = target.shape
-
-    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
-    num_warps = get_num_warps(BLOCK_SIZE)
+    BLOCK_SIZE = (
+        min(8192, triton.next_power_of_2(V))
+        if infer_device() == "xpu"
+        else min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+    )
+    num_warps = 32 if infer_device() == "xpu" else get_num_warps(BLOCK_SIZE)
 
     grid = (BT,)
 
@@ -185,9 +192,9 @@ class LigerKLDivLossFunction(torch.autograd.Function):
     Class implementing the forward and backward pass for the KL Divergence Loss using Triton, as defined by the following formula:
     ```python
     if log_target:
-        loss = target * (target.log() - input)
-    else:
         loss = target.exp() * (target - input)
+    else:
+        loss = target * (target.log() - input)
     ```,
     then the loss is reduced according to the `reduction` parameter.
     as defined in the PyTorch documentation: https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html

liger_kernel/ops/layer_norm.py

@@ -154,6 +154,11 @@ def layer_norm_forward(X, W, B, eps):
             f"must match weight size (W.shape[0]={W.shape[0]})"
         )
 
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args["grf_mode"] = "large"
+
     _layer_norm_forward_kernel[(n_rows,)](
         Y,
         Y.stride(0),
@@ -171,6 +176,7 @@ def layer_norm_forward(X, W, B, eps):
         eps,
         BLOCK_SIZE=BLOCK_SIZE,
         num_warps=num_warps,
+        **kernel_args,  # XPU-specific optimization
     )
     return Y.view(*shape), X, Mean, RSTD, BLOCK_SIZE, num_warps
 
@@ -185,7 +191,7 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
     if X.device.type == "cuda":
         sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
     elif X.device.type == "xpu":
-        sm_count = torch.xpu.get_device_properties(X.device).gpu_subslice_count
+        sm_count = torch.xpu.get_device_properties(X.device).gpu_eu_count
 
     DX = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
     _DW = torch.empty((sm_count, n_cols), dtype=W.dtype, device=W.device)
@@ -208,6 +214,12 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
         if X.dtype == torch.float16
         else tl.float32  # fallback to float32 for other types
     )
+
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args.update({"grf_mode": "large", "num_warps": 32, "num_stages": 4})
+
     _layer_norm_backward_kernel[grid](
         X,
         W,
@@ -227,6 +239,7 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
         rows_per_program,
         BLOCK_SIZE=BLOCK_SIZE,
         dtype=triton_dtype,
+        **kernel_args,  # XPU-specific optimization
     )
 
     DW = _DW.sum(dim=0).to(W.dtype)

liger_kernel/ops/rms_norm.py

@@ -223,6 +223,10 @@ def rms_norm_forward(X, W, eps, offset, casting_mode):
     # Check constraints.
     assert X.shape[1] == W.shape[0], "Incompatible hidden size dimension between tensor1.shape[1] and tensor2.shape[0]"
 
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args["grf_mode"] = "large"
     _rms_norm_forward_kernel[(n_rows,)](
         Y,
         Y.stride(0),
@@ -238,6 +242,7 @@ def rms_norm_forward(X, W, eps, offset, casting_mode):
         casting_mode,
         BLOCK_SIZE=BLOCK_SIZE,
         num_warps=num_warps,
+        **kernel_args,  # XPU-specific optimization
     )
     return Y.view(*shape), X, RSTD, BLOCK_SIZE, num_warps, casting_mode
 
@@ -252,7 +257,7 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
     if X.device.type == "cuda":
         sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
     elif X.device.type == "xpu":
-        sm_count = torch.xpu.get_device_properties(X.device).gpu_subslice_count
+        sm_count = torch.xpu.get_device_properties(X.device).gpu_eu_count
 
     # fp32 for numerical stability especially.
     _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
@@ -267,6 +272,11 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
     else:
         dX = torch.zeros_like(dY)
 
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args["grf_mode"] = "large"
+
     _rms_norm_backward_kernel[grid](
         dY,
         dY.stride(0),
@@ -288,6 +298,7 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
         casting_mode,
         BLOCK_SIZE=BLOCK_SIZE,
         num_warps=num_warps,
+        **kernel_args,  # XPU-specific optimization
     )
     dX = dX.view(*shape)
     dW = _dW.sum(dim=0).to(W.dtype)

liger_kernel/transformers/__init__.py

@@ -1,27 +1,145 @@
-from liger_kernel.transformers.auto_model import AutoLigerKernelForCausalLM  # noqa: F401
+import importlib
+
+from typing import TYPE_CHECKING
+
+# Always-safe imports (independent of 'transformers')
 from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss  # noqa: F401
+from liger_kernel.transformers.dyt import LigerDyT  # noqa: F401
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss  # noqa: F401
 from liger_kernel.transformers.fused_linear_jsd import LigerFusedLinearJSD  # noqa: F401
 from liger_kernel.transformers.geglu import LigerGEGLUMLP  # noqa: F401
 from liger_kernel.transformers.jsd import LigerJSD  # noqa: F401
 from liger_kernel.transformers.layer_norm import LigerLayerNorm  # noqa: F401
-from liger_kernel.transformers.monkey_patch import _apply_liger_kernel  # noqa: F401
-from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma2  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_granite  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mistral  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mixtral  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mllama  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_olmo2  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_phi3  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_5_vl  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_vl  # noqa: F401
 from liger_kernel.transformers.rms_norm import LigerRMSNorm  # noqa: F401
 from liger_kernel.transformers.rope import liger_rotary_pos_emb  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerSwiGLUMLP  # noqa: F401
 from liger_kernel.transformers.tvd import LigerTVDLoss  # noqa: F401
+
+# Static-only imports for IDEs and type checkers
+if TYPE_CHECKING:
+    from liger_kernel.transformers.auto_model import AutoLigerKernelForCausalLM  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import _apply_liger_kernel  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma2  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3_text  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_granite  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llava  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mistral  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mixtral  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mllama  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_olmo2  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_paligemma  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_phi3  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_5_vl  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_vl  # noqa: F401
+
+
+# Check if 'transformers' is installed
+try:
+    import transformers  # noqa: F401
+
+    _TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    _TRANSFORMERS_AVAILABLE = False
+
+
+def is_transformers_available() -> bool:
+    """
+    Returns True if the 'transformers' package is available.
+    Useful for conditional logic in downstream code.
+    """
+    return _TRANSFORMERS_AVAILABLE
+
+
+def __getattr__(name: str):
+    """
+    Handles lazy access to transformer-dependent attributes.
+    If 'transformers' is not installed, raises a user-friendly ImportError.
+    """
+    if not _TRANSFORMERS_AVAILABLE:
+        raise ImportError(
+            f"The attribute '{name}' requires the 'transformers' library, which is not installed.\n"
+            f"Please install it with `pip install transformers` to use this functionality."
+        )
+
+    if name == "AutoLigerKernelForCausalLM":
+        module = importlib.import_module("liger_kernel.transformers.auto_model")
+        return getattr(module, name)
+
+    monkey_patch_symbols = {
+        "_apply_liger_kernel",
+        "_apply_liger_kernel_to_instance",
+        "apply_liger_kernel_to_gemma",
+        "apply_liger_kernel_to_gemma2",
+        "apply_liger_kernel_to_gemma3",
+        "apply_liger_kernel_to_gemma3_text",
+        "apply_liger_kernel_to_granite",
+        "apply_liger_kernel_to_llama",
+        "apply_liger_kernel_to_llava",
+        "apply_liger_kernel_to_mistral",
+        "apply_liger_kernel_to_mixtral",
+        "apply_liger_kernel_to_mllama",
+        "apply_liger_kernel_to_olmo2",
+        "apply_liger_kernel_to_paligemma",
+        "apply_liger_kernel_to_phi3",
+        "apply_liger_kernel_to_qwen2",
+        "apply_liger_kernel_to_qwen2_5_vl",
+        "apply_liger_kernel_to_qwen2_vl",
+    }
+
+    if name in monkey_patch_symbols:
+        module = importlib.import_module("liger_kernel.transformers.monkey_patch")
+        return getattr(module, name)
+
+    raise AttributeError(f"module {__name__} has no attribute {name}")
+
+
+# Shared symbols in all environments
+__all__ = [
+    "is_transformers_available",
+    "LigerCrossEntropyLoss",
+    "LigerDyT",
+    "LigerFusedLinearCrossEntropyLoss",
+    "LigerFusedLinearJSD",
+    "LigerGEGLUMLP",
+    "LigerJSD",
+    "LigerLayerNorm",
+    "LigerRMSNorm",
+    "liger_rotary_pos_emb",
+    "LigerBlockSparseTop2MLP",
+    "LigerPhi3SwiGLUMLP",
+    "LigerSwiGLUMLP",
+    "LigerTVDLoss",
+]
+
+# Add transformer-dependent symbols only if available
+if _TRANSFORMERS_AVAILABLE:
+    __all__.extend(
+        [
+            "AutoLigerKernelForCausalLM",
+            "_apply_liger_kernel",
+            "_apply_liger_kernel_to_instance",
+            "apply_liger_kernel_to_gemma",
+            "apply_liger_kernel_to_gemma2",
+            "apply_liger_kernel_to_gemma3",
+            "apply_liger_kernel_to_gemma3_text",
+            "apply_liger_kernel_to_granite",
+            "apply_liger_kernel_to_llama",
+            "apply_liger_kernel_to_llava",
+            "apply_liger_kernel_to_mistral",
+            "apply_liger_kernel_to_mixtral",
+            "apply_liger_kernel_to_mllama",
+            "apply_liger_kernel_to_olmo2",
+            "apply_liger_kernel_to_paligemma",
+            "apply_liger_kernel_to_phi3",
+            "apply_liger_kernel_to_qwen2",
+            "apply_liger_kernel_to_qwen2_5_vl",
+            "apply_liger_kernel_to_qwen2_vl",
+        ]
+    )

liger_kernel/transformers/dyt.py

@@ -0,0 +1,20 @@
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops.dyt import LigerDyTFunction
+
+
+class LigerDyT(nn.Module):
+    def __init__(self, hidden_size, init_alpha=0.5):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.init_alpha = init_alpha
+        self.alpha = nn.Parameter(torch.ones(1) * init_alpha)
+        self.gamma = nn.Parameter(torch.ones(hidden_size))
+        self.beta = nn.Parameter(torch.zeros(hidden_size))
+
+    def forward(self, x):
+        return LigerDyTFunction.apply(x, self.alpha, self.gamma, self.beta)
+
+    def extra_repr(self):
+        return f"{self.hidden_size}, init_alpha={self.init_alpha}"

liger_kernel/transformers/functional.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
+from liger_kernel.ops.dyt import LigerDyTFunction
 from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
 from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
 from liger_kernel.ops.geglu import LigerGELUMulFunction
@@ -192,3 +193,7 @@ def liger_rope(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
 
 def liger_swiglu(a, b):
     return LigerSiLUMulFunction.apply(a, b)
+
+
+def liger_dyt(x, alpha, gamma, beta):
+    return LigerDyTFunction.apply(x, alpha, gamma, beta)

liger_kernel/transformers/gema3_rms.py

@@ -0,0 +1,8 @@
+from .rms_norm import LigerRMSNorm
+
+
+class LigerRMSNormForGemma3(LigerRMSNorm):
+    """Gemma3RMSNorm has a dim argument not hidden_size used in q_norm and k_norm."""
+
+    def __init__(self, dim, eps=0.000001, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False):
+        super().__init__(dim, eps, offset, casting_mode, init_fn, in_place)

liger_kernel/transformers/model/gemma.py

@@ -12,8 +12,10 @@ from transformers.models.gemma.modeling_gemma import _CONFIG_FOR_DOC
 from transformers.models.gemma.modeling_gemma import GEMMA_INPUTS_DOCSTRING
 from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
+from transformers.utils.deprecation import deprecate_kwarg
 
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 
 @add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
@@ -126,6 +128,7 @@ def lce_forward_deprecated(
     )
 
 
+@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
 @add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
 @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
 def lce_forward(
@@ -141,7 +144,7 @@ def lce_forward(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
-    num_logits_to_keep: int = 0,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -151,10 +154,12 @@ def lce_forward(
             config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
             (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
 
-        num_logits_to_keep (`int`, *optional*):
-            Calculate logits for the last `num_logits_to_keep` tokens. If `0`, calculate logits for all
+        logits_to_keep (`int` or `torch.Tensor`, *optional*):
+            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
             `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
             token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+            This is useful when using packed tensor format (single dimension for batch and sequence length).
 
     Returns:
 
@@ -200,24 +205,16 @@ def lce_forward(
     loss = None
     # if in training mode, don't materialize logits
     if self.training and (labels is not None):
-        # We do the same thing as ForCausalLMLoss but using Liger FLCE
-
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # flatten tokens
-        shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
-        shift_labels = shift_labels.view(-1)
-
-        reduction = "sum" if "num_items_in_batch" in loss_kwargs else "mean"
-        lce = LigerFusedLinearCrossEntropyLoss(reduction=reduction)
-
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
-        if reduction == "sum":
-            loss /= loss_kwargs["num_items_in_batch"]
-
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
     else:  # if in inference mode materialize logits
-        logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
         if labels is not None:
             loss = self.loss_function(
                 logits=logits,

liger_kernel/transformers/model/gemma2.py

@@ -13,8 +13,10 @@ from transformers.models.gemma2.modeling_gemma2 import _CONFIG_FOR_DOC
 from transformers.models.gemma2.modeling_gemma2 import GEMMA2_INPUTS_DOCSTRING
 from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
+from transformers.utils.deprecation import deprecate_kwarg
 
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 logger = logging.getLogger(__name__)
 
@@ -133,6 +135,7 @@ def lce_forward_deprecated(
     )
 
 
+@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
 @add_start_docstrings_to_model_forward(GEMMA2_INPUTS_DOCSTRING)
 @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
 def lce_forward(
@@ -148,7 +151,7 @@ def lce_forward(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
-    num_logits_to_keep: int = 0,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -158,10 +161,12 @@ def lce_forward(
             config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
             (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
 
-        num_logits_to_keep (`int`, *optional*):
-            Calculate logits for the last `num_logits_to_keep` tokens. If `0`, calculate logits for all
+        logits_to_keep (`int` or `torch.Tensor`, *optional*):
+            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
             `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
             token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+            This is useful when using packed tensor format (single dimension for batch and sequence length).
 
     Returns:
 
@@ -212,27 +217,18 @@ def lce_forward(
     loss = None
     # if in training mode, don't materialize logits
     if self.training and (labels is not None):
-        # We do the same thing as ForCausalLMLoss but using Liger FLCE
-
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # flatten tokens
-        shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
-        shift_labels = shift_labels.view(-1)
-
-        reduction = "sum" if "num_items_in_batch" in loss_kwargs else "mean"
-        lce = LigerFusedLinearCrossEntropyLoss(
-            softcap=self.config.final_logit_softcapping,
-            reduction=reduction,
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            final_logit_softcapping=self.config.final_logit_softcapping,
+            **loss_kwargs,
         )
 
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
-        if reduction == "sum":
-            loss /= loss_kwargs["num_items_in_batch"]
-
     else:  # if in inference mode materialize logits
-        logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
         if self.config.final_logit_softcapping is not None:
             logits = logits / self.config.final_logit_softcapping
             logits = torch.tanh(logits)