feat(flce): expose explicit token chunk sizes

Thread a validated chunk_size override through Triton, native CuTe, dispatch adapters, and Transformer wrappers while preserving backend defaults and legacy argument compatibility. Keep this layer independent of cuTile chunking and the three-way benchmark.

Co-authored-by: Copilot <[email protected]>

README.md

@@ -231,6 +231,47 @@ Ops without a CuTe DSL kernel transparently fall back to the default Triton kern
 
 The `cutedsl` extra also pulls in `apache-tvm-ffi`, which lets compiled kernels take PyTorch tensors directly rather than marshalling each one per call. It is optional — every kernel falls back to the marshalling launch without it — but short kernels are dominated by that per-call cost, so installing it is strongly recommended.
 
+### Fused Linear Cross Entropy `chunk_size`
+
+The fused-linear-cross-entropy path partitions the `B*T` token dimension into row
+chunks to bound the transient `chunk_size x V` logits buffer. By default each backend
+derives that chunk size from its own memory heuristic; pass an explicit `chunk_size` to
+override it. The same keyword is threaded through all three public entry points:
+
+```python
+import liger_kernel.functional as LF
+from liger_kernel.transformers.functional import liger_fused_linear_cross_entropy
+from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
+
+# top-level functional (keyword-only); always returns a 4-tuple
+loss, _z_loss, _token_accuracy, _predicted_tokens = LF.fused_linear_cross_entropy(x, weight, target, chunk_size=1024)
+
+# transformers functional
+loss = liger_fused_linear_cross_entropy(x, weight, target, chunk_size=1024)
+
+# nn.Module (set once at construction)
+loss_fn = LigerFusedLinearCrossEntropyLoss(chunk_size=1024)
+loss = loss_fn(weight, x, target)
+```
+
+Policy for an explicit override:
+
+- **Positive int or `None`.** `chunk_size` must be a positive Python `int`; `None` (the
+  default) keeps each backend's own heuristic. Zero, negative, `bool`, and non-int
+  values raise `ValueError`.
+- **No power-of-two rounding.** The value is honored exactly (e.g. `chunk_size=30` yields
+  a 30-row token loop), unlike the backends' default heuristics, which apply their own
+  backend-specific power-of-two policy (the shared Triton / shared CuTe-DSL orchestration
+  rounds up; the native SM100 CuTe DSL FLCE rounds down to a power of two and caps the chunk).
+- **Clamped to `B*T`.** A value larger than the token count is clamped so a single chunk
+  covers every row.
+- **Supported backends.** The shared Triton FLCE orchestration accepts `chunk_size` under both
+  of its public `impl=` selections -- `nvidia-triton` (Triton inner CE kernel) and `nvidia-cutedsl`
+  (the *same* shared orchestration with a CuTe DSL inner CE kernel, not a standalone native kernel).
+  The distinct, self-contained native SM100 CuTe DSL FLCE (selected via `LIGER_KERNEL_IMPL=cutedsl`)
+  also accepts it. The legacy cuTile FLCE backend does not yet advertise support and rejects an
+  explicit `chunk_size` (a later layer wires it through) rather than silently ignoring the override.
+
 ### Fused Scaled Cross Entropy
 
 `LigerFusedLinearScaledCrossEntropyFunction` is an additional per-token operator, not a replacement for the reduction-oriented Triton `LigerFusedLinearCrossEntropyFunction`. It takes `input[M, H]`, `weight[V, H]`, and `target[M]`, applies `logits / temperature`, and returns FP32 negative log-likelihood `[M]` plus optional differentiable vocabulary entropy `[M]` in the input dtype. Reductions remain in PyTorch, and rows whose target equals `ignore_index` contribute zero outputs and gradients.

src/liger_kernel/functional.py

@@ -678,6 +678,7 @@ def fused_linear_cross_entropy(
     return_token_accuracy: bool = False,
     return_predicted_tokens: bool = False,
     *,
+    chunk_size: Optional[int] = None,
     impl: Optional[str] = None,
     mode: Optional[str] = None,
     backend: Optional[str] = None,  # legacy back-compat alias for impl=
@@ -705,6 +706,10 @@ def fused_linear_cross_entropy(
         use_token_scaling: whether to scale each token's loss by its predicted probability.
         return_token_accuracy: if True, also return token-level accuracy.
         return_predicted_tokens: if True, also return predicted token indices.
+        chunk_size: optional explicit token-chunk size for the fused CE/GEMM loop. When
+            ``None`` (default) each backend uses its own memory heuristic. When set, it is
+            forwarded only to backends that advertise support (the Triton and shared
+            CuTe-DSL orchestration); an unsupported backend raises rather than ignoring it.
         impl: explicit implementation name (``"nvidia-cutedsl"``,
             ``"nvidia-triton"``). Legacy bare names accepted and translated.
         backend: Legacy back-compat alias for ``impl=``.
@@ -716,6 +721,9 @@ def fused_linear_cross_entropy(
     """
     if impl is None and backend is not None:
         impl = backend
+    # Only forward ``chunk_size`` when explicitly requested so legacy adapters that do not
+    # accept the kwarg (e.g. ascend-triton) keep working unchanged for the common None case.
+    extra_kwargs = {} if chunk_size is None else {"chunk_size": chunk_size}
     return dispatch(
         "fused_linear_cross_entropy",
         _input,
@@ -735,6 +743,7 @@ def fused_linear_cross_entropy(
         return_predicted_tokens,
         impl=impl,
         mode=mode,
+        **extra_kwargs,
     )
 
 

src/liger_kernel/ops/backends/_ascend/fused_linear_cross_entropy.py

@@ -42,9 +42,12 @@ def fused_linear_cross_entropy_ascend(
     return_predicted_tokens: bool = False,
     *,
     mode: Optional[str] = None,
+    chunk_size: Optional[int] = None,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
     if mode not in (None, "default"):
         raise ValueError(f"Ascend fused_linear_cross_entropy has only mode='default'; got mode={mode!r}.")
+    if chunk_size is not None:
+        raise ValueError("Ascend fused_linear_cross_entropy does not support an explicit chunk_size override.")
     return LigerFusedLinearCrossEntropyFunction.apply(
         _input,
         weight,

src/liger_kernel/ops/backends/_cutedsl/fused_linear_cross_entropy.py

@@ -70,6 +70,7 @@ def fused_linear_cross_entropy_cutedsl(
     return_predicted_tokens: bool = False,
     *,
     mode: Optional[str] = None,
+    chunk_size: Optional[int] = None,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
     """CuTe DSL fused_linear_cross_entropy dispatch entry point.
 
@@ -78,7 +79,8 @@ def fused_linear_cross_entropy_cutedsl(
     dispatcher — picking up the CuTe DSL kernel on Hopper+.
 
     ``mode`` is accepted only for API parity with the other backends; the only
-    valid value is ``"default"`` (or ``None``).
+    valid value is ``"default"`` (or ``None``). ``chunk_size`` optionally
+    overrides the token-chunk heuristic of the shared Triton orchestration.
     """
     if mode not in (None, "default"):
         raise ValueError(f"CuTe DSL fused_linear_cross_entropy has only mode='default'; got mode={mode!r}.")
@@ -100,4 +102,5 @@ def fused_linear_cross_entropy_cutedsl(
         return_predicted_tokens,
         "nvidia-cutedsl",
         mode,
+        chunk_size,
     )

src/liger_kernel/ops/backends/_triton/fused_linear_cross_entropy.py

@@ -61,10 +61,11 @@ def fused_linear_cross_entropy_triton(
     return_predicted_tokens: bool = False,
     *,
     mode: Optional[str] = None,
+    chunk_size: Optional[int] = None,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
     """Triton fused_linear_cross_entropy. ``mode`` is accepted for API parity
     and must be one of ``None``/``"default"``; anything else is rejected with a
-    clear error.
+    clear error. ``chunk_size`` optionally overrides the token-chunk heuristic.
     """
     if mode not in (None, "default"):
         raise ValueError(
@@ -89,4 +90,5 @@ def fused_linear_cross_entropy_triton(
         return_predicted_tokens,
         "nvidia-triton",
         mode,
+        chunk_size,
     )

src/liger_kernel/ops/cutedsl/ops/fused_linear_cross_entropy.py

@@ -10,6 +10,7 @@ from liger_kernel.ops.cutedsl.ops.cross_entropy import _launch_ce_fwd
 from liger_kernel.ops.utils import amp_custom_bwd
 from liger_kernel.ops.utils import amp_custom_fwd
 from liger_kernel.ops.utils import compare_version
+from liger_kernel.ops.utils import validate_flce_chunk_size
 from liger_kernel.utils import infer_device_arch
 
 _SUPPORTS_OUT_DTYPE = compare_version("torch", operator.ge, "2.8.0")
@@ -82,11 +83,13 @@ def _native_forward(
     return_token_accuracy,
     return_predicted_tokens,
     needs_grad,
+    chunk_size=None,
 ):
     """Compute every native output and gradient from one chunked logits pass."""
     BT, H = X.shape
     V = W.shape[0]
-    chunk = _forward_grad_chunk_size(BT)
+    explicit_chunk = validate_flce_chunk_size(chunk_size, BT)
+    chunk = explicit_chunk if explicit_chunk is not None else _forward_grad_chunk_size(BT)
     vector_width = 16 // X.element_size()
     storage_width = ((V + vector_width - 1) // vector_width) * vector_width
     logits_storage = torch.empty(chunk, storage_width, device=X.device, dtype=X.dtype)
@@ -247,6 +250,7 @@ def fused_linear_cross_entropy_forward(
     use_token_scaling=False,
     return_token_accuracy=False,
     return_predicted_tokens=False,
+    chunk_size=None,
 ):
     """Run the native SM100 CuTe DSL path."""
     ctx = _FlceState()
@@ -285,6 +289,10 @@ def fused_linear_cross_entropy_forward(
         "use_token_scaling": use_token_scaling,
         "return_token_accuracy": return_token_accuracy,
         "return_predicted_tokens": return_predicted_tokens,
+        # Persist the explicit chunk-size override so the repeated-backward recompute
+        # path (which re-invokes ``_native_forward(**native_kwargs)``) partitions the
+        # tokens identically to this forward pass. ``None`` keeps the default heuristic.
+        "chunk_size": chunk_size,
     }
     loss, z_out, acc_out, pred_out, gX, gW, gb = _native_forward(
         X,
@@ -349,7 +357,7 @@ def _flce_backward(ctx, grad_output):
             ctx._repeated_grad_bias = gb
         grads = (ctx._repeated_grad_input, ctx._repeated_grad_weight, ctx._repeated_grad_bias)
     gX, gW, gb = _scale_gradients(grads, grad_output, in_place=handoff)
-    return (gX, gW, None, gb, None, None, None, None, None, None, None, None, None, None, None)
+    return (gX, gW, None, gb, None, None, None, None, None, None, None, None, None, None, None, None, None, None)
 
 
 def fused_linear_cross_entropy_backward(ctx, grad_output):
@@ -361,6 +369,12 @@ def fused_linear_cross_entropy_backward(ctx, grad_output):
 
 
 class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
+    # The native SM100 CuTe DSL FLCE is a self-contained implementation: it does not
+    # dispatch its inner CE through the multi-backend registry, so ``ce_impl``/``ce_mode``
+    # are accepted only as API-parity placeholders (``None`` or the identity selection of
+    # this very impl). It does support an explicit token ``chunk_size`` override.
+    supports_chunk_size = True
+
     @staticmethod
     @amp_custom_fwd
     def forward(
@@ -380,7 +394,17 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
         use_token_scaling=False,
         return_token_accuracy=False,
         return_predicted_tokens=False,
+        ce_impl=None,
+        ce_mode=None,
+        chunk_size=None,
     ):
+        if ce_impl not in (None, "nvidia-cutedsl"):
+            raise ValueError(
+                "Native CuTe DSL FLCE does not dispatch an inner CE implementation; "
+                f"ce_impl must be None or 'nvidia-cutedsl'. Got: {ce_impl!r}"
+            )
+        if ce_mode not in (None, "default"):
+            raise ValueError(f"Native CuTe DSL FLCE only supports ce_mode=None or 'default'. Got: {ce_mode!r}")
         loss, z_loss, token_accuracy, predicted_tokens, state = fused_linear_cross_entropy_forward(
             _input,
             weight,
@@ -397,6 +421,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
             use_token_scaling,
             return_token_accuracy,
             return_predicted_tokens,
+            chunk_size=chunk_size,
         )
         if getattr(state, "saved_tensors", None) is not None:
             ctx.save_for_backward(*state.saved_tensors)

src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -14,6 +14,7 @@ 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.ops.utils import validate_flce_chunk_size
 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
@@ -61,6 +62,7 @@ def fused_linear_cross_entropy_forward(
     token_grad_output=None,
     compute_gradients=None,
     weight_requires_grad=None,
+    chunk_size=None,
 ):
     assert isinstance(return_z_loss, bool), f"return_z_loss must be True or False. Got: {return_z_loss}"
     assert isinstance(return_token_accuracy, bool), (
@@ -88,10 +90,17 @@ def fused_linear_cross_entropy_forward(
     BT, H = _input.shape
     V = weight.shape[0]
 
-    # widen the transient logits budget to C x BT x H (C=1 is a memory floor, not a perf target)
-    inc_factor = triton.cdiv(V, _CHUNK_MEM_CONST * H)
-    chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor))  # (BT + inc_factor - 1) // inc_factor
-    chunk_size = min(chunk_size, BT)  # a single chunk covers BT when the budget allows; never exceed BT
+    # An explicit ``chunk_size`` override wins over the default memory heuristic; it is
+    # validated (positive int, not bool) and clamped to BT without power-of-two rounding.
+    # ``None`` preserves the historical C=16 transient-budget heuristic below.
+    explicit_chunk_size = validate_flce_chunk_size(chunk_size, BT)
+    if explicit_chunk_size is not None:
+        chunk_size = explicit_chunk_size
+    else:
+        # widen the transient logits budget to C x BT x H (C=1 is a memory floor, not a perf target)
+        inc_factor = triton.cdiv(V, _CHUNK_MEM_CONST * H)
+        chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor))  # (BT + inc_factor - 1) // inc_factor
+        chunk_size = min(chunk_size, BT)  # a single chunk covers BT when the budget allows; never exceed BT
     num_chunks = triton.cdiv(BT, chunk_size)  # (BT + chunk_size - 1) // chunk_size
 
     grad_input = torch.zeros_like(_input, device=device)
@@ -413,6 +422,7 @@ def fused_linear_cross_entropy_backward(grad_output, grad_input, grad_weight, gr
 
 class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
     supports_inner_impl_dispatch = True
+    supports_chunk_size = True
 
     @staticmethod
     @amp_custom_fwd
@@ -435,6 +445,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
         return_predicted_tokens: bool = False,
         ce_impl=None,
         ce_mode=None,
+        chunk_size=None,
     ):
         """
         Fusing the last linear layer with cross-entropy loss
@@ -460,6 +471,9 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
             Default: False.
         return_token_accuracy (bool): When `return_token_accuracy` is `True`, computes and returns per-token accuracy without materializing logits. Default: `False`
         return_predicted_tokens (bool): When `return_predicted_tokens` is `True`, returns per-token predicted class indices (argmax) without materializing logits. Default: `False`
+        chunk_size (int): explicit override for the number of tokens processed per chunk. When `None`
+            (default) the memory heuristic (transient budget C=16) picks the chunk size. When set, it
+            must be a positive int; it is clamped to `B*T` and used verbatim (no power-of-two rounding).
         """
 
         # With reduction="none" the loss is per-token, so backward receives a
@@ -490,6 +504,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
                 return_predicted_tokens=return_predicted_tokens,
                 ce_impl=ce_impl,
                 ce_mode=ce_mode,
+                chunk_size=chunk_size,
                 compute_gradients=False if ctx.defer_grads else None,
             )
         )
@@ -509,6 +524,9 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
                 # reduction="none" path re-enters forward from backward to recompute grads.
                 ce_impl=ce_impl,
                 ce_mode=ce_mode,
+                # Persist the explicit chunk-size override so the reduction="none" backward
+                # recomputation partitions the tokens identically to the forward pass.
+                chunk_size=chunk_size,
             )
             ctx.weight_requires_grad = weight.requires_grad
         else:
@@ -571,4 +589,5 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
             None,  # return_predicted_tokens
             None,  # ce_impl
             None,  # ce_mode
+            None,  # chunk_size
         )

src/liger_kernel/ops/utils.py

@@ -74,6 +74,32 @@ def calculate_settings(n):
     return BLOCK_SIZE, num_warps
 
 
+def validate_flce_chunk_size(chunk_size, total_rows: int):
+    """Validate and clamp an explicit fused-linear-cross-entropy chunk-size override.
+
+    The fused linear cross-entropy (FLCE) implementations partition the ``BT`` token
+    dimension into row chunks to bound the transient ``chunk_size x V`` logits buffer.
+    Every backend derives a default ``chunk_size`` from a memory heuristic; this helper
+    validates and normalizes a caller-supplied *explicit* override so all backends share
+    one contract:
+
+    * ``None`` -> return ``None`` (the caller keeps its own default heuristic).
+    * otherwise ``chunk_size`` must be a positive Python ``int`` (``bool`` is rejected,
+      since ``True``/``False`` are almost always a mistake here).
+    * the value is clamped to ``total_rows`` (a single chunk can cover every row); it is
+      **not** rounded to a power of two — the explicit request is honored exactly.
+
+    Returns the validated (and possibly clamped) ``chunk_size``.
+    """
+    if chunk_size is None:
+        return None
+    if isinstance(chunk_size, bool) or not isinstance(chunk_size, int):
+        raise ValueError(f"chunk_size must be a positive int or None. Got: {chunk_size!r}")
+    if chunk_size <= 0:
+        raise ValueError(f"chunk_size must be a positive int or None. Got: {chunk_size}")
+    return min(chunk_size, total_rows)
+
+
 def compare_version(package: str, operator: Callable, target: str):
     try:
         pkg = importlib.import_module(package)

src/liger_kernel/transformers/functional.py

@@ -40,6 +40,40 @@ class CrossEntropyOutput:
     predicted_tokens: Optional[torch.Tensor] = None
 
 
+def build_flce_apply_args(flce_cls, base_args, chunk_size=None, ce_impl=None, ce_mode=None):
+    """Append the optional trailing FLCE arguments the target Function class supports.
+
+    ``base_args`` is the 15-tuple of positional args that *every* fused-linear-cross-entropy
+    ``autograd.Function`` accepts. The active class is backend-dependent (Triton, native
+    CuTe DSL, cuTile, Ascend, ...) because ``liger_kernel.ops`` swaps it based on
+    ``LIGER_KERNEL_IMPL``. We consult the class's capability metadata to decide what to
+    append so callers never pass unsupported positional args:
+
+    * ``supports_chunk_size`` -> append ``(ce_impl, ce_mode, chunk_size)``.
+    * ``supports_inner_impl_dispatch`` only -> append ``(ce_impl, ce_mode)``.
+    * legacy classes (e.g. Ascend / non-CUDA) -> append nothing.
+
+    A requested ``chunk_size`` (or ``ce_impl``/``ce_mode``) that the class cannot honor
+    raises a ``ValueError`` rather than being silently ignored.
+    """
+    supports_chunk = getattr(flce_cls, "supports_chunk_size", False)
+    supports_dispatch = getattr(flce_cls, "supports_inner_impl_dispatch", False)
+    if supports_chunk:
+        return base_args + (ce_impl, ce_mode, chunk_size)
+    if chunk_size is not None:
+        raise ValueError(
+            f"{getattr(flce_cls, '__name__', flce_cls)} does not support an explicit "
+            "chunk_size override on this backend."
+        )
+    if supports_dispatch:
+        return base_args + (ce_impl, ce_mode)
+    if ce_impl is not None or ce_mode is not None:
+        raise ValueError(
+            f"{getattr(flce_cls, '__name__', flce_cls)} does not support ce_impl/ce_mode selection on this backend."
+        )
+    return base_args
+
+
 # conform to the function signature in https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html
 # `weight` and `size_average` are placeholders and not implemented yet
 def liger_cross_entropy(
@@ -95,8 +129,9 @@ def liger_fused_linear_cross_entropy(
     use_token_scaling: bool = False,
     return_token_accuracy: bool = False,
     return_predicted_tokens: bool = False,
+    chunk_size: Optional[int] = None,
 ):
-    apply_args = (
+    base_args = (
         input,
         weight,
         target,
@@ -113,8 +148,7 @@ def liger_fused_linear_cross_entropy(
         return_token_accuracy,
         return_predicted_tokens,
     )
-    if getattr(LigerFusedLinearCrossEntropyFunction, "supports_inner_impl_dispatch", False):
-        apply_args += (None, None)
+    apply_args = build_flce_apply_args(LigerFusedLinearCrossEntropyFunction, base_args, chunk_size=chunk_size)
     loss, z_loss, token_accuracy, predicted_tokens = LigerFusedLinearCrossEntropyFunction.apply(*apply_args)
 
     if not return_z_loss and not return_token_accuracy and not return_predicted_tokens:

src/liger_kernel/transformers/fused_linear_cross_entropy.py

@@ -4,6 +4,7 @@ import torch
 
 from liger_kernel.ops import LigerFusedLinearCrossEntropyFunction
 from liger_kernel.transformers.functional import CrossEntropyOutput
+from liger_kernel.transformers.functional import build_flce_apply_args
 
 
 class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
@@ -20,6 +21,7 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         use_token_scaling: bool = False,
         return_token_accuracy: bool = False,
         return_predicted_tokens: bool = False,
+        chunk_size: Optional[int] = None,
     ):
         super().__init__()
         assert (label_smoothing >= 0) and (label_smoothing <= 1), (
@@ -42,9 +44,10 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         self.use_token_scaling = use_token_scaling
         self.return_token_accuracy = return_token_accuracy
         self.return_predicted_tokens = return_predicted_tokens
+        self.chunk_size = chunk_size
 
     def forward(self, lin_weight, _input, target, bias=None):
-        loss, z_loss, token_accuracy, predicted_tokens = LigerFusedLinearCrossEntropyFunction.apply(
+        base_args = (
             _input,
             lin_weight,
             target,
@@ -60,9 +63,12 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
             self.use_token_scaling,
             self.return_token_accuracy,
             self.return_predicted_tokens,
-            None,
-            None,
         )
+        # Only append the trailing ce_impl/ce_mode/chunk_size args the active backend
+        # Function advertises support for. The prior unconditional ``(None, None)`` broke
+        # backends (native CuTe DSL, cuTile, Ascend) whose Function takes only 15 args.
+        apply_args = build_flce_apply_args(LigerFusedLinearCrossEntropyFunction, base_args, chunk_size=self.chunk_size)
+        loss, z_loss, token_accuracy, predicted_tokens = LigerFusedLinearCrossEntropyFunction.apply(*apply_args)
         if not self.return_z_loss and not self.return_token_accuracy and not self.return_predicted_tokens:
             return loss
 

test/ops/test_flce_chunk_size.py

@@ -0,0 +1,592 @@
+"""Focused tests for the explicit FLCE ``chunk_size`` controls.
+
+In this layer only the backends that already ship an explicit token ``chunk_size``
+override participate: the shared Triton FLCE and the native SM100 CuTe DSL FLCE. The
+cuTile backend is intentionally left at its legacy 15-arg contract here (chunking lands
+in a later layer), so its only assertion below is the *negative* one — a non-``None``
+override is refused, never silently ignored.
+
+Covers:
+
+* the shared validation/clamp helper (``validate_flce_chunk_size``) — pure Python,
+* numeric parity of an explicit chunk vs. the default heuristic (Triton + native CuTe),
+* clamping (``chunk_size >= BT``) and non-power-of-two tails,
+* invalid overrides (zero / negative / bool / float) raising,
+* actual projection-GEMM row-geometry traces proving the explicit ``chunk_size`` really
+  partitions the token loop (Triton ``aten::mm`` shapes; native CuTe epilogue-GEMM rows),
+* the ``reduction="none"`` deferred-backward recomputation honoring the override, incl.
+  a non-uniform upstream gradient vs. an fp32 reference (Triton) and the native CuTe
+  repeated-backward recompute re-deriving the same chunking,
+* the env-facing wrappers (top-level + transformers functional + nn.Module),
+* the capability-gated dispatch helper refusing an override on unsupported backends,
+* the legacy cuTile Function rejecting an explicit override on this layer.
+
+Optional CUDA backends (native CuTe DSL SM100) are skipped when the current GPU / SDK
+does not satisfy them.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from torch.utils._python_dispatch import TorchDispatchMode
+
+from liger_kernel.ops.utils import validate_flce_chunk_size
+
+_HAS_CUDA = torch.cuda.is_available()
+
+
+def _cc():
+    return torch.cuda.get_device_capability() if _HAS_CUDA else (0, 0)
+
+
+def _native_cutedsl_available():
+    # Native CuTe DSL FLCE is SM100-only.
+    if not _HAS_CUDA or _cc() != (10, 0):
+        return False
+    try:
+        import cutlass.cute  # noqa: F401
+    except (ImportError, ModuleNotFoundError):
+        # Only a *missing* optional dependency skips the test; a broken but
+        # installed SDK must surface, not be silently hidden here.
+        return False
+    return True
+
+
+@pytest.fixture(autouse=True)
+def _require_cuda_for_ops_tests():
+    """Override the directory-level CUDA auto-skip for this mixed CPU/GPU file.
+
+    This module splits into pure-CPU groups (``validate_flce_chunk_size`` and the
+    dispatch-helper arity logic) and real-GPU groups. Every GPU group carries its
+    own explicit ``@pytest.mark.skipif`` guard, so the CPU groups must run even
+    when no CUDA device is visible. Yield without skipping instead of inheriting
+    the conftest ``autouse`` skip.
+    """
+    yield
+
+
+# ---------------------------------------------------------------------------
+# Pure-Python helper (no GPU required for the logic itself)
+# ---------------------------------------------------------------------------
+
+
+class TestValidateHelper:
+    def test_none_passthrough(self):
+        assert validate_flce_chunk_size(None, 128) is None
+
+    def test_positive_int_unchanged_when_below_total(self):
+        assert validate_flce_chunk_size(32, 128) == 32
+
+    def test_clamped_to_total_rows(self):
+        assert validate_flce_chunk_size(1024, 128) == 128
+        assert validate_flce_chunk_size(128, 128) == 128
+
+    def test_non_power_of_two_preserved_not_rounded(self):
+        # explicit request is honored exactly (no power-of-two rounding)
+        assert validate_flce_chunk_size(30, 100) == 30
+        assert validate_flce_chunk_size(97, 100) == 97
+
+    @pytest.mark.parametrize("bad", [0, -1, -128])
+    def test_non_positive_rejected(self, bad):
+        with pytest.raises(ValueError):
+            validate_flce_chunk_size(bad, 128)
+
+    @pytest.mark.parametrize("bad", [True, False])
+    def test_bool_rejected(self, bad):
+        with pytest.raises(ValueError):
+            validate_flce_chunk_size(bad, 128)
+
+    @pytest.mark.parametrize("bad", [3.5, 16.0, "16", 2.0])
+    def test_non_int_rejected(self, bad):
+        with pytest.raises(ValueError):
+            validate_flce_chunk_size(bad, 128)
+
+
+# ---------------------------------------------------------------------------
+# Shared references / helpers for the GPU parity tests
+# ---------------------------------------------------------------------------
+
+
+def _make_inputs(BT, H, V, dtype, seed=0, bias=False, device="cuda"):
+    torch.manual_seed(seed)
+    x = torch.randn(BT, H, device=device, dtype=dtype).requires_grad_(True)
+    w = (torch.randn(V, H, device=device, dtype=dtype) / (H**0.5)).detach().requires_grad_(True)
+    t = torch.randint(0, V, (BT,), device=device)
+    b = torch.randn(V, device=device, dtype=dtype).requires_grad_(True) if bias else None
+    return x, w, t, b
+
+
+def _torch_reference(x, w, t, bias=None, ignore_index=-100, reduction="mean"):
+    xr = x.detach().clone().float().requires_grad_(True)
+    wr = w.detach().clone().float().requires_grad_(True)
+    logits = xr @ wr.t()
+    if bias is not None:
+        logits = logits + bias.detach().float()
+    loss = F.cross_entropy(logits, t, ignore_index=ignore_index, reduction=reduction)
+    grad = torch.ones_like(loss) if reduction == "none" else None
+    loss.backward(grad)
+    return loss.detach(), xr.grad, wr.grad
+
+
+def _triton_flce():
+    from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
+
+    return LigerFusedLinearCrossEntropyFunction
+
+
+def _run_triton(x, w, t, bias=None, reduction="mean", chunk_size=None, accum_dtype=None):
+    Fn = _triton_flce()
+    loss, _, _, _ = Fn.apply(
+        x,
+        w,
+        t,
+        bias,
+        None,
+        -100,
+        0.0,
+        0.0,
+        reduction,
+        None,
+        False,
+        accum_dtype,
+        False,
+        False,
+        False,
+        None,
+        None,
+        chunk_size,
+    )
+    return loss
+
+
+# ---------------------------------------------------------------------------
+# Triton backend (universal — runs on any CUDA GPU)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not _HAS_CUDA, reason="CUDA required")
+class TestTritonChunkSize:
+    def test_parity_default_vs_explicit(self):
+        x, w, t, _ = _make_inputs(128, 64, 512, torch.float32, seed=1)
+        loss_def = _run_triton(x, w, t)
+        loss_def.backward()
+        gx_def, gw_def = x.grad.clone(), w.grad.clone()
+        x.grad = None
+        w.grad = None
+        loss_ex = _run_triton(x, w, t, chunk_size=16)
+        loss_ex.backward()
+        assert torch.allclose(loss_def, loss_ex, atol=1e-4, rtol=1e-4)
+        assert torch.allclose(gx_def, x.grad, atol=1e-3, rtol=1e-3)
+        assert torch.allclose(gw_def, w.grad, atol=1e-3, rtol=1e-3)
+
+    def test_parity_vs_torch_reference(self):
+        x, w, t, _ = _make_inputs(96, 48, 256, torch.float32, seed=2)
+        ref_loss, ref_gx, ref_gw = _torch_reference(x, w, t)
+        loss = _run_triton(x, w, t, chunk_size=30)  # non-power-of-two tail: 96 = 30*3 + 6
+        loss.backward()
+        assert torch.allclose(loss, ref_loss, atol=1e-3, rtol=1e-3)
+        assert torch.allclose(x.grad, ref_gx, atol=1e-3, rtol=1e-3)
+        assert torch.allclose(w.grad, ref_gw, atol=1e-3, rtol=1e-3)
+
+    def test_clamp_above_bt_matches_single_chunk(self):
+        x, w, t, _ = _make_inputs(64, 32, 128, torch.float32, seed=3)
+        loss_big = _run_triton(x, w, t, chunk_size=100000)  # clamped to BT=64
+        loss_big.backward()
+        gx_big = x.grad.clone()
+        x.grad = None
+        w.grad = None
+        loss_bt = _run_triton(x, w, t, chunk_size=64)
+        loss_bt.backward()
+        assert torch.allclose(loss_big, loss_bt, atol=1e-5, rtol=1e-5)
+        assert torch.allclose(gx_big, x.grad, atol=1e-5, rtol=1e-5)
+
+    @pytest.mark.parametrize("bad", [0, -1, True, 3.5])
+    def test_invalid_chunk_rejected(self, bad):
+        x, w, t, _ = _make_inputs(32, 16, 64, torch.float32, seed=4)
+        with pytest.raises(ValueError):
+            _run_triton(x, w, t, chunk_size=bad)
+
+    def test_reduction_none_recompute_honors_override(self):
+        # reduction="none" defers grads to backward, recomputing chunked logits.
+        # The explicit chunk_size must be persisted into that recomputation.
+        x, w, t, _ = _make_inputs(100, 40, 256, torch.float32, seed=5, bias=True)
+        ref_loss, ref_gx, ref_gw = _torch_reference(x, w, t, reduction="none")
+        Fn = _triton_flce()
+        loss, _, _, _ = Fn.apply(
+            x, w, t, None, None, -100, 0.0, 0.0, "none", None, False, None, False, False, False, None, None, 24
+        )
+        loss.sum().backward()  # upstream grad of ones -> matches reference reduction="none"
+        assert loss.shape == (100,)
+        assert torch.allclose(x.grad, ref_gx, atol=1e-3, rtol=1e-3)
+        assert torch.allclose(w.grad, ref_gw, atol=1e-3, rtol=1e-3)
+
+
+# ---------------------------------------------------------------------------
+# Native CuTe DSL backend (SM100 only)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not _native_cutedsl_available(), reason="native CuTe DSL FLCE requires SM100 + cutlass")
+class TestNativeCuTeChunkSize:
+    def _fn(self):
+        from liger_kernel.ops.cutedsl.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
+
+        return LigerFusedLinearCrossEntropyFunction
+
+    def _run(self, x, w, t, chunk_size=None, ce_impl=None, ce_mode=None):
+        Fn = self._fn()
+        loss, _, _, _ = Fn.apply(
+            x,
+            w,
+            t,
+            None,
+            None,
+            -100,
+            0.0,
+            0.0,
+            "mean",
+            None,
+            False,
+            torch.float32,
+            False,
+            False,
+            False,
+            ce_impl,
+            ce_mode,
+            chunk_size,
+        )
+        return loss
+
+    def test_supports_flag(self):
+        assert self._fn().supports_chunk_size is True
+
+    def test_parity_default_vs_explicit(self):
+        x, w, t, _ = _make_inputs(256, 256, 4096, torch.bfloat16, seed=6)
+        loss_def = self._run(x, w, t)
+        loss_def.backward()
+        gx_def, gw_def = x.grad.clone(), w.grad.clone()
+        x.grad = None
+        w.grad = None
+        loss_ex = self._run(x, w, t, chunk_size=64)
+        loss_ex.backward()
+        assert torch.allclose(loss_def, loss_ex, atol=2e-2, rtol=2e-2)
+        assert torch.allclose(gx_def, x.grad, atol=1e-1, rtol=1e-1)
+        assert torch.allclose(gw_def, w.grad, atol=1e-1, rtol=1e-1)
+
+    def test_non_power_of_two_tail(self):
+        x, w, t, _ = _make_inputs(200, 256, 4096, torch.bfloat16, seed=7)
+        loss = self._run(x, w, t, chunk_size=48)  # 200 = 48*4 + 8
+        loss.backward()
+        assert torch.isfinite(loss)
+
+    @pytest.mark.parametrize("bad", [0, -3, True, 2.5])
+    def test_invalid_chunk_rejected(self, bad):
+        x, w, t, _ = _make_inputs(64, 256, 4096, torch.bfloat16, seed=8)
+        with pytest.raises(ValueError):
+            self._run(x, w, t, chunk_size=bad)
+
+    def test_ce_impl_mode_placeholders(self):
+        x, w, t, _ = _make_inputs(64, 256, 4096, torch.bfloat16, seed=9)
+        # self-identity / default accepted
+        self._run(x, w, t, ce_impl="nvidia-cutedsl", ce_mode="default", chunk_size=32)
+        with pytest.raises(ValueError):
+            self._run(x, w, t, ce_impl="nvidia-triton")
+        with pytest.raises(ValueError):
+            self._run(x, w, t, ce_mode="static_persistent")
+
+
+# ---------------------------------------------------------------------------
+# cuTile backend (SM90 + SM100) — chunking is owned by a later layer. Here we only
+# assert the *negative* contract: this layer leaves cuTile at its legacy 15-arg
+# Function, so an explicit override must be refused, never silently ignored.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(_cc() not in ((9, 0), (10, 0)), reason="cuTile FLCE requires SM90/SM100")
+class TestCuTileLegacyRejectsChunk:
+    def _fn(self):
+        pytest.importorskip("cuda.tile")
+        from liger_kernel.ops.cutile.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
+
+        return LigerFusedLinearCrossEntropyFunction
+
+    def test_does_not_advertise_chunk_support(self):
+        # cuTile stays on the legacy contract in this layer.
+        assert getattr(self._fn(), "supports_chunk_size", False) is False
+
+    def test_explicit_override_refused_by_dispatch_helper(self):
+        from liger_kernel.transformers.functional import build_flce_apply_args
+
+        base = tuple(range(15))
+        with pytest.raises(ValueError):
+            build_flce_apply_args(self._fn(), base, chunk_size=64)
+
+
+# ---------------------------------------------------------------------------
+# Projection-GEMM row-geometry traces — prove the explicit ``chunk_size`` really
+# partitions the token loop (C is used, not ignored). Adapted (minimal) from the
+# combined branch's test/ops/test_flce_benchmark.py; no benchmark harness here.
+# ---------------------------------------------------------------------------
+
+
+def _expected_chunk_rows(tokens, chunk):
+    """Row count of each projection chunk after clamping ``chunk`` to ``tokens``."""
+    chunk = min(chunk, tokens)
+    rows, start = [], 0
+    while start < tokens:
+        rows.append(min(chunk, tokens - start))
+        start += chunk
+    return rows
+
+
+def _flce_inputs_cuda(tokens, hidden, vocab, seed=0, dtype=torch.bfloat16):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    x = torch.randn(tokens, hidden, device="cuda", dtype=dtype, generator=g).requires_grad_(True)
+    w = (torch.randn(vocab, hidden, device="cuda", dtype=dtype, generator=g) / (hidden**0.5)).detach()
+    w.requires_grad_(True)
+    target = torch.randint(0, vocab, (tokens,), device="cuda", generator=g)
+    return x, w, target
+
+
+def _apply_direct(fn, x, w, target, chunk_size, reduction="mean"):
+    loss, _, _, _ = fn.apply(
+        x,
+        w,
+        target,
+        None,
+        None,
+        -100,
+        0.0,
+        0.0,
+        reduction,
+        None,
+        False,
+        torch.float32,
+        False,
+        False,
+        False,
+        None,
+        None,
+        chunk_size,
+    )
+    return loss
+
+
+class _ProjectionMMRecorder(TorchDispatchMode):
+    """Record the LHS/RHS shapes of every ``aten::mm`` the forward issues."""
+
+    def __init__(self):
+        self.mms = []
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        kwargs = kwargs or {}
+        out = func(*args, **kwargs)
+        if str(func).startswith("aten.mm"):
+            a, b = args[0], args[1]
+            self.mms.append((tuple(a.shape), tuple(b.shape)))
+        return out
+
+
+_TOKENS, _HIDDEN, _VOCAB, _CHUNK = 101, 256, 4096, 33
+
+
+@pytest.mark.skipif(not _HAS_CUDA, reason="CUDA required")
+class TestTritonProjectionGeometry:
+    def _fn(self):
+        from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
+
+        return LigerFusedLinearCrossEntropyFunction
+
+    def test_projection_row_geometry_matches_chunking(self):
+        # Non-power-of-two override: 101 tokens / chunk 33 -> [33, 33, 33, 2].
+        x, w, target = _flce_inputs_cuda(_TOKENS, _HIDDEN, _VOCAB, seed=20)
+        with _ProjectionMMRecorder() as rec:
+            _apply_direct(self._fn(), x, w, target, _CHUNK)
+        # Projection GEMM is x_chunk @ w.T -> lhs [rows, H], rhs [H, V].
+        proj = [lhs for lhs, rhs in rec.mms if rhs == (_HIDDEN, _VOCAB)]
+        rows = [lhs[0] for lhs in proj]
+        assert all(lhs[1] == _HIDDEN for lhs in proj)
+        assert rows == _expected_chunk_rows(_TOKENS, _CHUNK)
+
+    def test_clamp_above_tokens_is_single_chunk(self):
+        x, w, target = _flce_inputs_cuda(_TOKENS, _HIDDEN, _VOCAB, seed=21)
+        with _ProjectionMMRecorder() as rec:
+            _apply_direct(self._fn(), x, w, target, _TOKENS * 5)  # clamps to N
+        rows = [lhs[0] for lhs, rhs in rec.mms if rhs == (_HIDDEN, _VOCAB)]
+        assert rows == [_TOKENS]
+
+    def test_reduction_none_backward_matches_fp32_ref_nonuniform_upstream(self):
+        # reduction="none" defers grads to backward, recomputing chunked logits with
+        # the persisted override; a non-uniform upstream grad exercises the full path.
+        tokens, hidden, vocab = 100, 40, 256
+        torch.manual_seed(22)
+        x = torch.randn(tokens, hidden, device="cuda", dtype=torch.float32, requires_grad=True)
+        w = (torch.randn(vocab, hidden, device="cuda", dtype=torch.float32) / (hidden**0.5)).detach()
+        w.requires_grad_(True)
+        target = torch.randint(0, vocab, (tokens,), device="cuda")
+        upstream = torch.rand(tokens, device="cuda", dtype=torch.float32) + 0.1  # nonuniform per-token
+
+        xr = x.detach().clone().requires_grad_(True)
+        wr = w.detach().clone().requires_grad_(True)
+        ref = F.cross_entropy(xr @ wr.t(), target, reduction="none")
+        ref.backward(upstream)
+
+        loss = _apply_direct(self._fn(), x, w, target, 24, reduction="none")
+        assert loss.shape == (tokens,)
+        loss.backward(upstream)
+        assert torch.allclose(x.grad, xr.grad, atol=1e-3, rtol=1e-3)
+        assert torch.allclose(w.grad, wr.grad, atol=1e-3, rtol=1e-3)
+
+
+@pytest.mark.skipif(not _native_cutedsl_available(), reason="native CuTe DSL FLCE requires SM100 + cutlass")
+class TestNativeCuTeProjectionGeometry:
+    def _fn(self):
+        from liger_kernel.ops.cutedsl.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
+
+        return LigerFusedLinearCrossEntropyFunction
+
+    def test_projection_epilogue_gemm_row_geometry(self, monkeypatch):
+        fn = self._fn()
+        import liger_kernel.ops.cutedsl.ops.fused_linear_cross_entropy as mod
+
+        real = mod.run_epilogue_gemm
+        rows = []
+
+        def _spy(Xc, W, dlogits, epilogue, *a, **k):
+            rows.append(dlogits.shape[0])
+            return real(Xc, W, dlogits, epilogue, *a, **k)
+
+        monkeypatch.setattr(mod, "run_epilogue_gemm", _spy)
+        x, w, target = _flce_inputs_cuda(_TOKENS, _HIDDEN, _VOCAB, seed=24)
+        _apply_direct(fn, x, w, target, _CHUNK)
+        assert rows == _expected_chunk_rows(_TOKENS, _CHUNK)
+
+    def test_repeated_backward_recompute_persists_chunk_override(self, monkeypatch):
+        fn = self._fn()
+        import liger_kernel.ops.cutedsl.ops.fused_linear_cross_entropy as mod
+
+        real = mod.run_epilogue_gemm
+        rows = []
+
+        def _spy(Xc, W, dlogits, epilogue, *a, **k):
+            rows.append(dlogits.shape[0])
+            return real(Xc, W, dlogits, epilogue, *a, **k)
+
+        monkeypatch.setattr(mod, "run_epilogue_gemm", _spy)
+        x, w, target = _flce_inputs_cuda(_TOKENS, _HIDDEN, _VOCAB, seed=25)
+
+        expected = _expected_chunk_rows(_TOKENS, _CHUNK)
+        loss = _apply_direct(fn, x, w, target, _CHUNK)
+        rows.clear()
+        loss.backward(retain_graph=True)  # 1st: consumes native fwd grads, no recompute
+        assert rows == []
+        loss.backward(retain_graph=True)  # 2nd: recompute -> same chunk override
+        assert rows == expected
+        rows.clear()
+        loss.backward(retain_graph=True)  # 3rd: cached -> no recompute
+        assert rows == []
+
+
+# ---------------------------------------------------------------------------
+# Capability-gated dispatch helper (pure Python)
+# ---------------------------------------------------------------------------
+
+
+class TestDispatchHelper:
+    def _helper(self):
+        from liger_kernel.transformers.functional import build_flce_apply_args
+
+        return build_flce_apply_args
+
+    def test_appends_triple_for_chunk_capable(self):
+        build = self._helper()
+
+        class Cap:
+            supports_chunk_size = True
+            supports_inner_impl_dispatch = True
+
+        base = tuple(range(15))
+        out = build(Cap, base, chunk_size=7)
+        assert out == base + (None, None, 7)
+
+    def test_appends_pair_for_dispatch_only(self):
+        build = self._helper()
+
+        class Cap:
+            supports_inner_impl_dispatch = True
+
+        base = tuple(range(15))
+        assert build(Cap, base) == base + (None, None)
+
+    def test_legacy_class_gets_no_extra_args(self):
+        build = self._helper()
+
+        class Legacy:
+            pass
+
+        base = tuple(range(15))
+        assert build(Legacy, base) == base
+
+    def test_chunk_on_unsupported_raises(self):
+        build = self._helper()
+
+        class Legacy:
+            supports_inner_impl_dispatch = True  # dispatch but no chunk support
+
+        base = tuple(range(15))
+        with pytest.raises(ValueError):
+            build(Legacy, base, chunk_size=8)
+
+    def test_ce_impl_on_legacy_raises(self):
+        build = self._helper()
+
+        class Legacy:
+            pass
+
+        base = tuple(range(15))
+        with pytest.raises(ValueError):
+            build(Legacy, base, ce_impl="nvidia-triton")
+
+
+# ---------------------------------------------------------------------------
+# Env-facing wrappers (top-level functional, transformers functional + nn.Module)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not _HAS_CUDA, reason="CUDA required")
+class TestEnvWrappers:
+    def test_top_level_functional_chunk(self):
+        import liger_kernel.functional as LF
+
+        x, w, t, _ = _make_inputs(128, 64, 512, torch.float32, seed=11)
+        ref_loss, _, _ = _torch_reference(x, w, t)
+        out = LF.fused_linear_cross_entropy(x, w, t, chunk_size=20)
+        assert torch.allclose(out[0], ref_loss, atol=1e-3, rtol=1e-3)
+
+    def test_transformers_functional_chunk(self):
+        from liger_kernel.transformers.functional import liger_fused_linear_cross_entropy
+
+        x, w, t, _ = _make_inputs(128, 64, 512, torch.float32, seed=12)
+        ref_loss, _, _ = _torch_reference(x, w, t)
+        loss = liger_fused_linear_cross_entropy(x, w, t, chunk_size=20)
+        assert torch.allclose(loss, ref_loss, atol=1e-3, rtol=1e-3)
+
+    def test_nn_module_chunk(self):
+        from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+
+        x, w, t, _ = _make_inputs(128, 64, 512, torch.float32, seed=13)
+        ref_loss, ref_gx, ref_gw = _torch_reference(x, w, t)
+        loss = LigerFusedLinearCrossEntropyLoss(chunk_size=20)(w, x, t)
+        loss.backward()
+        assert torch.allclose(loss, ref_loss, atol=1e-3, rtol=1e-3)
+        assert torch.allclose(x.grad, ref_gx, atol=1e-3, rtol=1e-3)
+        assert torch.allclose(w.grad, ref_gw, atol=1e-3, rtol=1e-3)
+
+    @pytest.mark.parametrize("bad", [0, -1, True])
+    def test_top_level_invalid_chunk_rejected(self, bad):
+        import liger_kernel.functional as LF
+
+        x, w, t, _ = _make_inputs(64, 32, 128, torch.float32, seed=14)
+        with pytest.raises(ValueError):
+            LF.fused_linear_cross_entropy(x, w, t, chunk_size=bad)