Keep Triton FLCE C=1 by default and use C=8 for GPT-OSS (#1472)

## Summary

- keep the public Triton fused linear cross entropy chunk-memory default
at `C=1`
- add an explicit per-call `chunk_mem_const` override while preserving
backward compatibility for existing autograd callers
- use `C=8` only in the GPT-OSS patched training path, where C=1 causes
a deterministic BF16 convergence mismatch
- keep every other model and direct/public FLCE caller at C=1
- retain the reusable-logits, direct-dX, and direct-dW optimizations
from #1468
- lock the default and override chunk geometry in regression coverage

## Why the override is scoped

The exact mini GPT-OSS convergence case fails with C=1 and passes with
C=4, C=8, and C=16. However, setting C=8 globally causes the BF16
Qwen3.5-MoE convergence case to fail. The model-scoped override
preserves both trajectories:

- public/default FLCE: C=1
- GPT-OSS patched loss: C=8
- fused linear KL: unchanged independently tuned policy

## H200 memory and latency

BF16, `BT=8192`, `H=4096`, `V=128256`:

| Policy | Chunk size | Chunks | Time | Incremental peak allocation |
|---|---:|---:|---:|---:|
| Public default C=1 | 256 | 32 | 53.039 ms | 1128.66 MiB |
| GPT-OSS override C=8 | 2048 | 4 | 37.748 ms | 1568.04 MiB |
| Previous global C=16 | 4096 | 2 | 35.074 ms | 2068.04 MiB |

The default therefore retains the minimum memory budget, while GPT-OSS
uses the smallest requested release override with stable convergence.

## Testing done

On H200 with PyTorch 2.9.1.6+cu129 and Triton 3.6.0:

- **144 passed**: FLCE ops suite
- **144 passed**: FLCE ops suite with `LIGER_KERNEL_STRICT=1`
- **141 passed**: transformer-facing FLCE suite
- **71 passed, 3 dependency-gated skips**: FP32 + BF16 text-model
convergence suites
- exact BF16 mini GPT-OSS convergence passes with the scoped C=8
override
- BF16 mini Qwen3.5-MoE convergence passes because its default remains
C=1
- Ruff lint/format and `git diff --check` pass

The broader six-file convergence attempt also exposed four PaliGemma
multimodal failures in the ad hoc Transformers 5.15 environment. Those
configurations explicitly set `fused_linear_cross_entropy=False`, so
they are unrelated to this change.

No owned B300 pod was available for an additional Blackwell run.

---------

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

src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -21,21 +21,10 @@ from liger_kernel.utils import infer_device
 # The optimal maximum block size depends on your hardware, your kernel, and your dtype
 MAX_FUSED_SIZE = 2048 if infer_device() == "npu" else 65536 // 2
 
-# Memory-budget constant for the token-chunk geometry below. The historical
-# formula (inc_factor = cdiv(V, H), i.e. budget C=1) sizes chunks to a memory
-# FLOOR of ~BT x H transient logits, which at LLM vocab shapes forces many tiny
-# chunks and a launch-bound wall. At llama_3 shapes (V=128256, H=4096) C=1 gives
-# inc_factor=32 -> chunk_size=BT/32 (256 rows @ BT=8192) = 32 loop iterations.
-# Widening the transient budget to C x BT x H collapses the loop. Measured
-# Blackwell B200 (full fwd+bwd, bf16): V=128256/H=4096 BT=8192 123.4ms (C=1) ->
-# 24.7ms (C=16) = 5.0x; BT=4096 119.3 -> 15.1 = 7.9x; BT=1024 116.9 -> 9.07 =
-# 12.9x; small V=32000 BT=8192 11.0 -> 5.4 = 2.0x (flat for C>=8). C=16 is the
-# knee (gains flatten beyond it at BT>=4096). Numerics are partition-invariant
-# per row: loss is bit-identical across C; input-grad differs only at the ~1e-5
-# GEMM reduction-order level; weight-grad differs only in addmm chunk-accumulation
-# order (same class as the fused_linear_jsd chunk-memory idiom). Same constant on
-# all devices.
-_CHUNK_MEM_CONST = 16
+# Public callers default to the historical minimum transient-logits budget:
+# approximately BT x H elements regardless of vocabulary size. Integrations
+# with stricter reduction-order requirements may explicitly request more.
+_CHUNK_MEM_CONST = 1
 _TORCH_VERSION = Version(torch.__version__.split("+")[0])
 _ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0")
 
@@ -80,6 +69,7 @@ def fused_linear_cross_entropy_forward(
     token_grad_output=None,
     compute_gradients=None,
     weight_requires_grad=None,
+    chunk_mem_const: int = _CHUNK_MEM_CONST,
 ):
     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,6 +78,8 @@ def fused_linear_cross_entropy_forward(
     assert isinstance(return_predicted_tokens, bool), (
         f"return_predicted_tokens must be True or False. Got: {return_predicted_tokens}"
     )
+    if isinstance(chunk_mem_const, bool) or not isinstance(chunk_mem_const, int) or chunk_mem_const < 1:
+        raise ValueError(f"chunk_mem_const must be a positive integer. Got: {chunk_mem_const!r}")
     device = _input.device
     # ``compute_gradients`` / ``weight_requires_grad`` let backward re-enter this
     # function with *detached* saved tensors (whose ``.requires_grad`` is False) and
@@ -107,8 +99,8 @@ 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)
+    # Bound transient logits to C x BT x H.
+    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
@@ -494,6 +486,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
         return_predicted_tokens: bool = False,
         ce_impl=None,
         ce_mode=None,
+        chunk_mem_const: int = _CHUNK_MEM_CONST,
     ):
         """
         Fusing the last linear layer with cross-entropy loss
@@ -519,6 +512,8 @@ 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_mem_const (int): transient-logits budget multiplier. Defaults to
+            `1`; larger values trade memory for fewer chunk reductions.
         """
 
         # With reduction="none" the loss is per-token, so backward receives a
@@ -550,6 +545,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
                 ce_impl=ce_impl,
                 ce_mode=ce_mode,
                 compute_gradients=False if ctx.defer_grads else None,
+                chunk_mem_const=chunk_mem_const,
             )
         )
 
@@ -568,6 +564,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
                 # reduction="none" path re-enters forward from backward to recompute grads.
                 ce_impl=ce_impl,
                 ce_mode=ce_mode,
+                chunk_mem_const=chunk_mem_const,
             )
             ctx.weight_requires_grad = weight.requires_grad
         else:
@@ -612,7 +609,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
             grad_input, grad_weight, grad_bias = fused_linear_cross_entropy_backward(
                 grad_output, grad_input, grad_weight, grad_bias
             )
-        return (
+        gradients = (
             grad_input,
             grad_weight,
             None,
@@ -630,4 +627,6 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
             None,  # return_predicted_tokens
             None,  # ce_impl
             None,  # ce_mode
+            None,  # chunk_mem_const
         )
+        return gradients[: len(ctx.needs_input_grad)]

src/liger_kernel/ops/fused_linear_kl_div.py

@@ -17,8 +17,8 @@ MAX_FUSED_SIZE = 4096 if infer_device() == "xpu" else 65536 // 2
 # Memory-budget constant for the token-chunk geometry below. The historical
 # formula (inc_factor = cdiv(V, H), i.e. budget C=1) sizes chunks to a memory
 # FLOOR of ~BT x H transient logits, which at LLM vocab shapes forces many tiny
-# chunks and a launch-bound wall. Same constant and rationale as
-# fused_linear_cross_entropy.py; see the measured numbers in its comment.
+# chunks and a launch-bound wall. Fused KL retains its independently tuned
+# C=16 policy; fused linear cross entropy defaults to the C=1 floor.
 _CHUNK_MEM_CONST = 16
 _TORCH_VERSION = Version(torch.__version__.split("+")[0])
 _ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0")

src/liger_kernel/transformers/functional.py

@@ -96,6 +96,7 @@ def liger_fused_linear_cross_entropy(
     use_token_scaling: bool = False,
     return_token_accuracy: bool = False,
     return_predicted_tokens: bool = False,
+    chunk_mem_const: int = 1,
 ):
     apply_args = (
         input,
@@ -116,6 +117,7 @@ def liger_fused_linear_cross_entropy(
     )
     if getattr(LigerFusedLinearCrossEntropyFunction, "supports_inner_impl_dispatch", False):
         apply_args += (None, None)
+    apply_args += (chunk_mem_const,)
     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/model/gpt_oss.py

@@ -11,6 +11,8 @@ from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result
 from liger_kernel.transformers.model.output_classes import LigerMoeCausalLMOutputWithPast
 
+_GPT_OSS_FLCE_CHUNK_MEM_CONST = 8
+
 
 def lce_forward(
     self,
@@ -138,6 +140,7 @@ def lce_forward(
     output_hidden_states = (
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
+    flce_chunk_mem_const = kwargs.pop("chunk_mem_const", _GPT_OSS_FLCE_CHUNK_MEM_CONST)
 
     # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
     outputs: MoeModelOutputWithPast = self.model(
@@ -175,6 +178,7 @@ def lce_forward(
             labels=labels,
             shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
+            chunk_mem_const=flce_chunk_mem_const,
             **kwargs,
         )
         loss, _, token_accuracy, predicted_tokens = unpack_cross_entropy_result(result)

test/ops/test_fused_linear_cross_entropy.py

@@ -273,12 +273,8 @@ def test_fused_linear_cross_entropy_propagates_backend_to_inner_ce(monkeypatch,
 # fixture, and the accumulation path itself requires a CUDA SM80+ device.
 #
 # The production token-chunk loop sizes chunks from the ``_CHUNK_MEM_CONST``
-# memory budget (``inc_factor = cdiv(V, C * H)``). At the tiny shapes used here
-# the default ``C=16`` collapses to a single chunk, which would never exercise
-# the per-chunk grad_weight accumulation these tests target. The class-local
-# autouse ``_force_multichunk`` fixture pins that budget knob to ``C=1`` *only
-# for these migrated tests* so the real production loop splits these shapes into
-# multiple chunks -- the cross-backend tests above keep the production ``C=16``.
+# memory budget (``inc_factor = cdiv(V, C * H)``). The public default is C=1,
+# so the small shapes below naturally exercise the real multi-chunk path.
 # ===========================================================================
 
 # ---------------------------------------------------------------------------
@@ -324,16 +320,20 @@ requires_nvidia_cuda = pytest.mark.skipif(
 )
 
 
-def _production_chunk_size(BT: int, V: int, H: int) -> int:
-    """The chunk size the production loop will pick for these dims, using the
-    *current* module-level ``_CHUNK_MEM_CONST`` (pinned to 1 by the class-local
-    autouse fixture). Mirrors the arithmetic in
-    ``fused_linear_cross_entropy_forward``."""
-    inc_factor = triton.cdiv(V, flce_ops._CHUNK_MEM_CONST * H)
+def _production_chunk_size(BT: int, V: int, H: int, chunk_mem_const: int | None = None) -> int:
+    """Mirror the production token-chunk geometry for the given dimensions."""
+    chunk_mem_const = flce_ops._CHUNK_MEM_CONST if chunk_mem_const is None else chunk_mem_const
+    inc_factor = triton.cdiv(V, chunk_mem_const * H)
     chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor))
     return min(chunk_size, BT)
 
 
+def test_triton_flce_chunk_memory_budget():
+    assert flce_ops._CHUNK_MEM_CONST == 1
+    assert _production_chunk_size(BT=8192, V=128256, H=4096) == 256
+    assert _production_chunk_size(BT=8192, V=128256, H=4096, chunk_mem_const=8) == 2048
+
+
 # ---------------------------------------------------------------------------
 # Dispatch recorder
 # ---------------------------------------------------------------------------
@@ -597,24 +597,7 @@ def _mk_inputs(
 
 
 class TestFusedLinearCrossEntropyAddmm:
-    """Regression suite for chunk buffers and the ``grad_weight`` addmm path.
-
-    Grouped in a class so the ``_force_multichunk`` autouse fixture (which pins
-    ``_CHUNK_MEM_CONST=1``) is scoped to *only* these migrated tests and never
-    perturbs the production ``C=16`` used by the cross-backend tests above.
-    """
-
-    @pytest.fixture(autouse=True)
-    def _force_multichunk(self, monkeypatch):
-        """Pin the existing ``_CHUNK_MEM_CONST`` memory budget to 1 for each test.
-
-        This is the *production* chunk-sizing knob (``inc_factor = cdiv(V, C*H)``),
-        not a test-only chunk API. At ``C=1`` the small shapes below split into
-        several token chunks, so the genuine per-chunk grad_weight accumulation
-        loop runs more than once -- which is exactly what these regression tests
-        cover. Being class-local, it does not affect any other test in the module.
-        """
-        monkeypatch.setattr(flce_ops, "_CHUNK_MEM_CONST", 1)
+    """Regression suite for chunk buffers and dW accumulation."""
 
     # -----------------------------------------------------------------------
     # 1. Dispatch: the low-precision default path uses addmm(out=) and no param upcast