bench(flce): add configurable chunk-memory sweep (#1479)

## Summary

- Add an opt-in benchmark sweep for FLCE `chunk_mem_const` values.
- Encode each chunk budget in the provider name so CSV speed and memory
rows remain distinguishable.
- Allow benchmark callers to configure Triton `do_bench` warmup and
repetition windows.
- Use longer measurement windows only when the chunk sweep is enabled.
- Reject combining the Triton chunk sweep with
`LIGER_KERNEL_IMPL=cutile`.
- Preserve existing benchmark providers, timings, and production FLCE
defaults when the sweep is disabled.

## Details

The shared benchmark repetition window can be shorter than a single FLCE
iteration for large-vocabulary shapes. In that case,
`triton.testing.do_bench` can effectively collect one sample and report
identical p20, p50, and p80 values.

Enable the sweep with:

```bash
LIGER_FLCE_BENCH_CHUNK_MEM_CONSTS="1,2,4,8,16" \
LIGER_FLCE_BENCH_REP_MS="5000" \
LIGER_BENCH_TARGET="flce_chunk_sweep" \
python benchmark/scripts/benchmark_fused_linear_cross_entropy.py \
  --model llama_3_8b \
  --overwrite
```

The resulting providers are recorded as `liger-chunk-c1`,
`liger-chunk-c2`, and so on.

Without `LIGER_FLCE_BENCH_CHUNK_MEM_CONSTS`, the script retains the
existing provider set and the shared 25 ms warmup / 10 ms repetition
defaults.

This is a benchmark-only change. It does not modify the production chunk
policy or recommend a hardware-wide default. Local SM89 measurements
showed that the speed/memory trade-off varies with token count and
vocabulary size.

This complements #1439: that PR exposes runtime chunk-size control,
while this PR provides reproducible tooling for evaluating chunk-memory
policies.

## Testing Done

- Hardware Type: NVIDIA GeForce RTX 4070 Laptop GPU (SM89), BF16, WSL2
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test` to ensure correctness
- The single-process suite was terminated by the local WSL2 host at 91%
and 93% on two attempts, in different transformer test files, consistent
with cumulative host-memory pressure.
  - `pytest -q test/ops/test_fused_linear_cross_entropy.py`: 144 passed.
- `pytest -q test/transformers/test_rms_norm.py`: 64 passed, 8 expected
XFAIL.
- [ ] run `make test-convergence` to ensure convergence
- Not completed locally. This PR changes benchmark tooling only and does
not modify production kernel behavior.

Representative benchmark validation:

- comma-separated chunk values are preserved in provider order;
- duplicate values are removed;
- invalid and non-positive values are rejected;
- a configurable repetition window produces non-degenerate p20/p50/p80
values for long-running C=1 paths;
- ascending and descending large-vocabulary sweeps showed less than 1.1%
maximum median order difference;
- default execution remains unchanged when the sweep environment
variables are absent.

Co-authored-by: Arup De <[email protected]>

benchmark/scripts/benchmark_fused_linear_cross_entropy.py

@@ -1,3 +1,14 @@
+"""Benchmark fused linear cross entropy implementations.
+
+Set ``LIGER_FLCE_BENCH_CHUNK_MEM_CONSTS`` to a comma-separated list (for
+example, ``1,2,4,8,16``) to compare Triton FLCE chunk-memory budgets. Provider
+names encode the selected value so speed and memory rows remain distinguishable
+in the shared benchmark CSV. Chunk sweeps use longer warmup and repetition
+windows because a single FLCE iteration can exceed the shared benchmark
+default's entire measurement budget. Override the repetition window with
+``LIGER_FLCE_BENCH_REP_MS`` for especially long-running shapes.
+"""
+
 import os
 
 import torch
@@ -12,11 +23,59 @@ from utils import build_speed_bench_fn
 from utils import parse_benchmark_script_args
 from utils import run_benchmarks
 
+from liger_kernel.transformers.functional import liger_fused_linear_cross_entropy
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
 from liger_kernel.utils import infer_device
 
 device = infer_device()
 
+_CHUNK_SWEEP_ENV = "LIGER_FLCE_BENCH_CHUNK_MEM_CONSTS"
+_CHUNK_SWEEP_REP_ENV = "LIGER_FLCE_BENCH_REP_MS"
+_CHUNK_PROVIDER_PREFIX = "liger-chunk-c"
+_CHUNK_SWEEP_WARMUP_MS = 100
+_CHUNK_SWEEP_REP_MS = 400
+
+
+def _parse_chunk_mem_consts(raw_value: str) -> tuple[int, ...]:
+    """Parse an ordered, comma-separated list of positive chunk budgets."""
+    if not raw_value.strip():
+        return ()
+
+    values = []
+    for raw_item in raw_value.split(","):
+        item = raw_item.strip()
+        if not item:
+            raise ValueError(f"{_CHUNK_SWEEP_ENV} must be a comma-separated list of positive integers")
+        try:
+            value = int(item)
+        except ValueError as exc:
+            raise ValueError(f"{_CHUNK_SWEEP_ENV} contains a non-integer value: {item!r}") from exc
+        if value < 1:
+            raise ValueError(f"{_CHUNK_SWEEP_ENV} values must be positive. Got: {value}")
+        if value not in values:
+            values.append(value)
+    return tuple(values)
+
+
+def _chunk_mem_const_from_provider(provider: str) -> int | None:
+    """Decode a chunk budget from a benchmark-only provider name."""
+    if not provider.startswith(_CHUNK_PROVIDER_PREFIX):
+        return None
+    return int(provider.removeprefix(_CHUNK_PROVIDER_PREFIX))
+
+
+def _parse_positive_int(raw_value: str, env_name: str, default: int) -> int:
+    """Parse a positive integer environment override."""
+    if not raw_value.strip():
+        return default
+    try:
+        value = int(raw_value)
+    except ValueError as exc:
+        raise ValueError(f"{env_name} must be a positive integer. Got: {raw_value!r}") from exc
+    if value < 1:
+        raise ValueError(f"{env_name} must be positive. Got: {value}")
+    return value
+
 
 class TorchLMHeadCE(torch.nn.Module):
     """Ground truth implementation of the linear fused with torch based cross entropy loss.
@@ -38,14 +97,35 @@ class TorchLMHeadCE(torch.nn.Module):
 
 
 class LigerLMHeadCE(torch.nn.Module):
-    def __init__(self, H: int, V: int, dtype: torch.dtype, ignore_index: int = -100, accum_dtype=None):
+    def __init__(
+        self,
+        H: int,
+        V: int,
+        dtype: torch.dtype,
+        ignore_index: int = -100,
+        accum_dtype=None,
+        chunk_mem_const: int | None = None,
+    ):
         super().__init__()
         self.lin = torch.nn.Linear(in_features=H, out_features=V, bias=False, dtype=dtype)
+        self.ignore_index = ignore_index
+        self.accum_dtype = accum_dtype
+        self.chunk_mem_const = chunk_mem_const
         self.ce_loss = LigerFusedLinearCrossEntropyLoss(
             ignore_index=ignore_index, reduction="mean", accum_dtype=accum_dtype
         )
 
     def forward(self, x, y):
+        if self.chunk_mem_const is not None:
+            return liger_fused_linear_cross_entropy(
+                input=x,
+                weight=self.lin.weight,
+                target=y,
+                ignore_index=self.ignore_index,
+                reduction="mean",
+                accum_dtype=self.accum_dtype,
+                chunk_mem_const=self.chunk_mem_const,
+            )
         return self.ce_loss(self.lin.weight, x, y)
 
 
@@ -67,7 +147,10 @@ def setup_fused_linear_cross_entropy(input: SingleBenchmarkRunInput):
     _input = torch.randn(BT, H, requires_grad=True, dtype=dtype, device=device)
     target = torch.randint(V, (BT, 1), dtype=torch.long, device=device).squeeze(1)
 
-    if input.kernel_provider == "liger":
+    chunk_mem_const = _chunk_mem_const_from_provider(input.kernel_provider)
+    if chunk_mem_const is not None:
+        lm_head_ce = LigerLMHeadCE(H=H, V=V, dtype=dtype, chunk_mem_const=chunk_mem_const).to(device)
+    elif input.kernel_provider == "liger":
         lm_head_ce = LigerLMHeadCE(H=H, V=V, dtype=dtype).to(device)
     elif input.kernel_provider == "liger-fp32-accum":
         lm_head_ce = LigerLMHeadCE(H=H, V=V, dtype=dtype, accum_dtype=torch.float32).to(device)
@@ -79,6 +162,16 @@ def setup_fused_linear_cross_entropy(input: SingleBenchmarkRunInput):
 if __name__ == "__main__":
     args = parse_benchmark_script_args()
     cutile_backend = os.environ.get("LIGER_KERNEL_IMPL", "").strip().lower() == "cutile"
+    chunk_mem_consts = _parse_chunk_mem_consts(os.environ.get(_CHUNK_SWEEP_ENV, ""))
+    if cutile_backend and chunk_mem_consts:
+        raise ValueError(
+            f"{_CHUNK_SWEEP_ENV} benchmarks Triton chunk budgets and cannot be combined with LIGER_KERNEL_IMPL=cutile"
+        )
+    chunk_sweep_rep_ms = _parse_positive_int(
+        os.environ.get(_CHUNK_SWEEP_REP_ENV, ""),
+        _CHUNK_SWEEP_REP_ENV,
+        _CHUNK_SWEEP_REP_MS,
+    )
 
     if args.sweep_mode == "model_config":
         common_configs = build_model_config_sweep(
@@ -112,12 +205,22 @@ if __name__ == "__main__":
             overwrite=args.overwrite,
         )
 
-    common_configs["kernel_providers"] = (
-        ["torch", "liger"] if cutile_backend else ["torch", "liger", "liger-fp32-accum"]
-    )
+    if chunk_mem_consts:
+        common_configs["kernel_providers"] = [
+            "torch",
+            *[f"{_CHUNK_PROVIDER_PREFIX}{chunk_mem_const}" for chunk_mem_const in chunk_mem_consts],
+        ]
+    else:
+        common_configs["kernel_providers"] = (
+            ["torch", "liger"] if cutile_backend else ["torch", "liger", "liger-fp32-accum"]
+        )
 
     run_benchmarks(
-        bench_test_fn=build_speed_bench_fn(setup_fused_linear_cross_entropy),
+        bench_test_fn=build_speed_bench_fn(
+            setup_fused_linear_cross_entropy,
+            warmup=_CHUNK_SWEEP_WARMUP_MS if chunk_mem_consts else 25,
+            rep=chunk_sweep_rep_ms if chunk_mem_consts else 10,
+        ),
         kernel_operation_modes=["forward", "full"] if cutile_backend else ["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",

benchmark/scripts/utils.py

@@ -119,6 +119,7 @@ def run_speed_benchmark(
     mode: str,
     input_tensors: List[torch.Tensor],
     rep: int = 10,
+    warmup: int = 25,
 ) -> "SingleBenchmarkRunOutput":
     """Measure execution speed for forward, backward, or full (fwd+bwd).
 
@@ -133,6 +134,7 @@ def run_speed_benchmark(
         ms_50, ms_20, ms_80 = triton.testing.do_bench(
             fwd_fn,
             grad_to_none=input_tensors,
+            warmup=warmup,
             rep=rep,
             quantiles=QUANTILES,
         )
@@ -142,6 +144,7 @@ def run_speed_benchmark(
         ms_50, ms_20, ms_80 = triton.testing.do_bench(
             lambda: y.backward(do, retain_graph=True),
             grad_to_none=input_tensors,
+            warmup=warmup,
             rep=rep,
             quantiles=QUANTILES,
         )
@@ -154,6 +157,7 @@ def run_speed_benchmark(
         ms_50, ms_20, ms_80 = triton.testing.do_bench(
             full,
             grad_to_none=input_tensors,
+            warmup=warmup,
             rep=rep,
             quantiles=QUANTILES,
         )
@@ -165,6 +169,7 @@ def run_speed_benchmark(
 
         ms_50, ms_20, ms_80 = triton.testing.do_bench(
             no_grad_forward,
+            warmup=warmup,
             rep=rep,
             quantiles=QUANTILES,
         )
@@ -210,10 +215,18 @@ def default_forward_fn(*setup_out):
 def build_speed_bench_fn(
     setup_fn: Callable[["SingleBenchmarkRunInput"], Any],
     forward_fn: Callable[..., torch.Tensor] = default_forward_fn,
+    warmup: int = 25,
+    rep: int = 10,
 ) -> Callable:
     def bench_speed(input: "SingleBenchmarkRunInput") -> SingleBenchmarkRunOutput:
         setup_out = setup_fn(input)
-        return run_speed_benchmark(lambda: forward_fn(*setup_out), input.kernel_operation_mode, [setup_out[0]])
+        return run_speed_benchmark(
+            lambda: forward_fn(*setup_out),
+            input.kernel_operation_mode,
+            [setup_out[0]],
+            warmup=warmup,
+            rep=rep,
+        )
 
     return bench_speed