feat(flce): expose direct chunk-size control via LIGER_FLCE_CHUNK_SIZE

[jin-five]

Summary

#1414 widened the FLCE token-chunk memory budget to C x BT x H with a hardcoded _CHUNK_MEM_CONST = 16 — a clear latency win, but the budget is linear in BT, so the transient [chunk_size, V] logits slab keeps growing on long sequences (V=262144, H=5376: 1.07 GiB at BT=8k, 2.1 GiB at 16k, 4.3 GiB at 32k), beyond the BT <= 8192 envelope measured in #1414.

On memory-tight runs that slab is a real binding transient on main today via the existing accum_dtype=torch.float32 path. Gemma-4 31B (V=262144, H=5376), DeepSpeed ZeRO-3, no offload, 8x H100 80 GB:

• seq 8k: with the default geometry the training run OOMs during step 5, 3 out of 3 attempts — allocator pressure at saturation, not the slab alone exceeding capacity (the same config survives with expandable_segments:True); LIGER_FLCE_CHUNK_SIZE=256 completes deterministically on the identical config, seed, and data — the ~0.9 GiB smaller slab is enough to move the run off that edge.
• seq 12k: the default geometry OOMs directly on the slab allocation (Tried to allocate 2.00 GiB at logits_chunk = _input_chunk @ weight.t()); a 256-row chunk clears it.

A small chunk is not universally better — the default is 3x faster where the memory exists (table below). The problem is that the slab is the only FLCE transient that grows with BT while being completely non-tunable. This PR exposes chunk_size directly via LIGER_FLCE_CHUNK_SIZE (rows; floored to a power of two, clamped to BT per call; read once at import time). Unset — the default — leaves the geometry identical to main. Set below the heuristic's choice it doubles as a BT-invariant bound on the slab (chunk_size x V x itemsize bytes); set above, it trades a larger slab for fewer loop iterations. num_chunks stays derived: a fixed chunk count would keep the slab BT-dependent, which is the original problem.

Details

Measured trade-off on H100 80GB (full fwd+bwd, bf16, default accumulation path, today's main; the 256-row geometry):

shape / BT	default	chunk_size=256	delta
V=262144 H=5376, BT=8192	146.9 ms / 11.58 GB	442.6 ms / 10.71 GB	default is 3.0x faster
V=262144 H=5376, BT=32768	462.7 ms / 14.83 GB	1735.3 ms / 10.95 GB	256-row saves 3.9 GB
V=128256 H=4096, BT=32768	163.2 ms / 9.18 GB	640.8 ms / 4.23 GB	256-row saves 5.0 GB

(peaks are transients over the persistent weights/grads, via torch.cuda.max_memory_allocated; note the 256-row peak is flat in BT — that is the point of the knob. Larger values, e.g. 1024, interpolate the trade-off.)

Testing Done

• Hardware Type: H100-80GB (CUDA 12.8, PyTorch 2.10, Triton 3.6)
• test/transformers/test_fused_linear_cross_entropy.py: 141 passed with the variable unset, =256, and =100 (the last normalizes to 64, exercising the power-of-two floor); unset path verified geometry-identical to main
• End-to-end paired run (31B, seq 8k, ZeRO-3, accum_dtype=fp32, same seed, 128 steps): default geometry OOMs during step 5 (3 out of 3 attempts); LIGER_FLCE_CHUNK_SIZE=256 completes (40.11 s/step, allocated peak 75.48 GiB vs 76.36 GiB default). The nominal headroom is misleading: non-PyTorch usage (CUDA context / NCCL / kernel modules) is ~2.6 GiB, so the observed usable allocation budget is roughly ~77 GiB, and the failing run is already at allocator saturation (a 2.62 GiB transient request with 0.75 GiB reserved-but-unallocated). The 0.88 GiB smaller slab moves the configuration off that edge
• run make test to ensure correctness — verified by paired same-window runs: this branch and unmodified main produce identical pass/fail/skip counts and identical failure sets in this environment (failures are environmental: corporate proxy blocks some model/asset downloads)
• run make checkstyle to ensure code style — ruff check clean; ruff format reports the same 13 pre-existing files on unmodified main
• run make test-convergence to ensure convergence — no new failures vs unmodified main (shared failures are the same environmental set)

Test logs (H100-80GB, CUDA 12.8, PyTorch 2.10, Triton 3.6)

FLCE unit suite — test/transformers/test_fused_linear_cross_entropy.py, run with the variable unset, LIGER_FLCE_CHUNK_SIZE=256, and =100 (identical result each time):

======================= 141 passed, 3 warnings in 4.24s ========================

Full test suite — paired same-window runs, this branch vs unmodified main (identical failure-name sets; all failures environmental — corporate proxy blocks some model/asset downloads):

branch: = 301 failed, 5148 passed, 1136 skipped, 3 xfailed =
main:   = 301 failed, 5148 passed, 1136 skipped, 3 xfailed =

make test-convergence — same paired protocol: 8 shared failures on both (falcon_h1 / paligemma, environmental); one borderline-tolerance test (qwen3_moe bf16) flips pass/fail on reruns of both the branch and unmodified main:

branch: = 9 failed, 132 passed, 36 skipped, 1 xpassed =
main:   = 8 failed, 133 passed, 36 skipped, 1 xpassed =   (qwen3_moe flips on both: 1 passed / 1 failed across two reruns each)

make checkstyle — ruff check clean; ruff format reports only the same 13 pre-existing files as unmodified main. The changed file itself:

$ ruff check src/liger_kernel/ops/fused_linear_cross_entropy.py
All checks passed!
$ ruff format --check src/liger_kernel/ops/fused_linear_cross_entropy.py
1 file already formatted

[hiwuhgds-pixel]

I agree that making chunking tunable is useful, but I’m not sure C is the right knob to expose.

C is only a proxy used to derive chunk_size and num_chunks. Because of cdiv and next_power_of_2, many different C values end up producing the same geometry.

For example, with V=262144, H=5376, BT=8192:

C	num_chunks	chunk_size
1	32	256
2–3	16	512
4–6	8	1024
7–16	4	2048

It may be cleaner to expose chunk_size or num_chunks directly, since those are the values that actually determine execution.

[jin-five]

@hiwuhgds-pixel
Thanks! You're right. I verified the collapse: at V=262144, H=5376, BT=8192, C=1..16 yields only four distinct geometries (chunk_size 256 / 512 / 1024 / 2048), with C=7..16 all mapping to 2048. I've reworked the branch accordingly (PR title/body updated): the C knob is gone, and LIGER_FLCE_CHUNK_SIZE now selects the chunk_size directly, after power-of-two flooring and a per-call clamp to BT, replacing the heuristic. Set below the heuristic's choice, it doubles as a BT-invariant bound on the [chunk_size, V] slab (chunk_size x V x itemsize bytes, the memory-tight case in the PR body); set above, it trades a larger slab for fewer loop iterations. Unset keeps the geometry identical to main.

I left num_chunks alone: a fixed chunk count keeps the slab growing with BT, which was the original problem.

Re-validated on the new variable: the FLCE unit suite passes with it unset, =256, and =100 (the last normalizes to 64, exercising the power-of-two floor), and the paired 31B seq-8k ZeRO-3 run from the PR body completes with LIGER_FLCE_CHUNK_SIZE=256 (40.11 s/step, allocated peak 75.48 GiB) where the default geometry OOMs during step 5, 3 out of 3 attempts.

[vaibhavjindal]

Hi @jin-five and @hiwuhgds-pixel, thanks for the PR and discussing this issue with long context training. I agree that we need a way to adjust the chunk_size on user demand while also having a default setting. Also, I feel like we should have this parameter as a part of the API instead of exposing it as an environment variable.

Also, we need to make sure the default behavior is good, and might want to have different defaults for different GPU types(example. H100 has about 80GB memory whereas B300 has about 280GB). I will discuss this further with my team at Linkedin and decide on the default behavior.

I am happy to contribute the change to your PR as a maintainer edit if that's okay with you, or leave the implementation to you if you'd prefer.

cc @kolehma8 @yueyiming2009 @arde171

[jin-five]

Hi @vaibhavjindal, thanks for taking a look!

Agreed on both points. Exposing chunk_size as part of the API (rather than an env var) makes sense, and GPU-dependent defaults sound reasonable since the memory headroom differs so much across devices.

Yes, please feel free to push the change directly to this PR as a maintainer edit. "Allow edits by maintainers" is enabled on the branch. Happy to help with testing or follow-ups once the default behavior is decided.

[jin-five]

@vaibhavjindal I checked the new release and saw that it includes the fused_linear_cross_entropy.py changes from #1472. Luckily the default _CHUNK_MEM_CONST is back to what it was before #1414, so it should give me back the tiny headroom I was missing.

After checking the latest release, I rebased this PR onto the recent main(19d92e7) to resolve a conflict on the _CHUNK_MEM_CONST block. And squashed my two commits into one.

Re-verified on H100 (torch 2.10.0+cu128, triton 3.6.0):

• LIGER_FLCE_CHUNK_SIZE unset: test/ops + test/transformers FLCE suites pass (285/285).
• LIGER_FLCE_CHUNK_SIZE=256 and =100: the 141 transformer-facing tests pass, but 68 tests in the new TestFusedLinearCrossEntropyAddmm class fail. That is expected rather than a kernel issue: the class mirrors the default geometry via _production_chunk_size and relies on small shapes being multi-chunk, so a global override collapses them into a single chunk. If we keep an env-style override, I'd add an autouse fixture there pinning _CHUNK_SIZE_OVERRIDE = 0, in the same spirit as the removed _force_multichunk fixture(not done yet).

Given that #1472 already resolves my original OOM and adds chunk_mem_const to the API as you suggested, I'd like to check the direction before going further. What #1472 doesn't cover is the point @hiwuhgds-pixel raised above: C is a proxy, and many C values collapse to the same chunk_size after cdiv + power-of-two rounding, while chunk_size itself is what determines the [chunk_size, V] slab and the loop count. The GPU-dependent default you mentioned is also still open.

I'll wait for the team's decision on the API shape and the default policy, and I'm happy to rework this PR accordingly (API argument instead of the env var, GPU-dependent default, or whatever direction you land on).