mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
7a12e092579dba20fa02388bf80a2c39cbc9010a
13
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
|
7a12e09257 |
[cuTile] Add v3 cuTile backends for 6 Liger ops (#1321)
Add cuTile (nvidia-cuda-tile) backends, dispatched via LIGER_KERNEL_IMPL=cutile, for eight ops — each drops into the existing LigerXxxFunction autograd interface so the backend swap is transparent: wiglu (SiLUMul), dyt, softmax, poly_norm, fused_add_rms_norm, grpo_loss (token- and sequence-level importance sampling, GRPO/CISPO/SAPO/VESPO) Correctness is exercised by the existing transformer suites run under LIGER_KERNEL_IMPL=cutile (which routes each op to its cuTile implementation), and test_cutile_backend.py verifies that routing. Benchmarks: Triton-vs-cuTile speed+memory data on NVIDIA B200 for all eight ops <img width="1000" height="600" alt="image" src="https://github.com/user-attachments/assets/fc534f41-51a5-4cd5-b02c-fc53930c5256" /> <img width="1000" height="600" alt="image" src="https://github.com/user-attachments/assets/0a3b7412-1503-47ad-a355-769b11574b53" /> <img width="1000" height="600" alt="image" src="https://github.com/user-attachments/assets/52a22e48-0982-4754-885f-9c957c3bc4ea" /> <img width="1000" height="600" alt="image" src="https://github.com/user-attachments/assets/0ab1102f-d74c-4d49-bf47-febc2db407f5" /> <img width="1000" height="600" alt="image" src="https://github.com/user-attachments/assets/6c6512fc-61af-46e3-ab68-63e3727ba7df" /> <img width="1000" height="600" alt="image" src="https://github.com/user-attachments/assets/8a912ce7-5fb4-4cc4-a761-fcd19fd68902" /> ## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
|
6470a51d3a |
Add SM100 CuTe DSL fused linear cross entropy (#1339)
## Summary
Adds a reusable CuTe DSL SM100 GEMM and uses it to implement fused
linear cross
entropy (FLCE) for NVIDIA B100/B200/GB200.
The shared GEMM provides:
- aligned and row-padded output views for bounded FLCE workspaces;
- a single-fragment custom epilogue entry point for future fused
operators;
- shared TMA/UMMA pipelines, persistent scheduling, TMEM management,
stream/device handling, compile serialization, and TVM-FFI/DLPack launch
paths.
FLCE computes each bounded logits chunk once, transforms that storage
into
`dlogits` in place with the CuTe DSL cross-entropy kernel, and
immediately forms
`grad_input`, `grad_weight`, and optional `grad_bias`. The first
autograd
backward hands off those precomputed gradients without recomputing
logits or
allocating a duplicate parameter-sized gradient.
At the Llama 3-sized reference shape (`BT=8192`, `H=4096`, `V=128256`,
BF16),
full forward + backward is **6.96x faster** than upstream Liger Triton
FLCE and
uses **2.17x less peak allocated memory** on B200.
## Implementation
- Native execution requires exact compute capability `(10, 0)` and
FP16/BF16
operands.
- The architecture gate uses the input tensor's device index and
`liger_kernel.utils.infer_device_arch`.
- The native path uses one bounded 512/1024-row workspace; a full `(BT,
V)`
logits tensor is never materialized.
- Cross entropy overwrites each logits chunk with normalized `dlogits`,
which
is consumed immediately by the input-, weight-, and bias-gradient
projections.
- Hidden dimensions not aligned to the native K tile are zero-padded
internally
and sliced back after backward.
- Odd vocabulary sizes use an aligned physical workspace stride while
the
logical vocabulary size excludes padding from softmax mass, label
smoothing,
and argmax.
- The single native path supports bias, class weights, label smoothing,
z-loss,
softcap, token scaling, token accuracy, predicted tokens, `mean`/`sum`
reduction, partial/frozen-gradient inputs, and configurable gradient
accumulation dtype.
- This implementation is native-only: non-SM100 hardware, non-FP16/BF16
operands, and `reduction="none"` raise an explicit error. Frontend
fallback
and cross-backend routing are intentionally outside this PR.
- Retained-graph repeated backward recomputes only after the first
precomputed
gradient handoff has been consumed.
## Scope
- `src/liger_kernel/ops/cutedsl/ops/_sm100_gemm.py`
- `src/liger_kernel/ops/cutedsl/ops/cross_entropy.py`
- `src/liger_kernel/ops/cutedsl/ops/fused_linear_cross_entropy.py`
- `src/liger_kernel/ops/cutedsl/ops/utils.py`
- `src/liger_kernel/ops/cutedsl/ops/__init__.py`
- `test/cutedsl/test_fused_linear_cross_entropy.py`
- `test/cutedsl/test_sm100_gemm.py`
- `benchmark/scripts/run_cutedsl_compare.py`
- `benchmark/scripts/utils.py`
- `benchmark/README.md`
`setup.py`, upstream Triton FLCE, the standard FLCE benchmark
implementation,
and SwiGLU are unchanged. This PR adds no Quack dependency and uses the
existing
`nvidia-cutlass-dsl` optional dependency.
## Correctness and validation
- FLCE suite: **236 passed**, with one environment-gated dispatch test
skipped
in the default run; that dispatch test passes explicitly with
`LIGER_KERNEL_IMPL=cutedsl`
- Generic SM100 GEMM epilogue suite: **4 passed**
- Mixed BF16-weight/FP32-bias gradients: passed across token- and
vocabulary-oriented, weighted and unweighted backward paths
- TVM-FFI-disabled DLPack launch fallback: passed
- Parallel first-compilation processes: passed
- Ruff check and format: passed
- Independent final code review: no high-confidence blockers
Coverage includes FP16/BF16, irregular and production shapes, aligned
and
predicated-store tails, odd vocabulary sizes, both scalar reductions,
ignored
targets, bias, class weights, label smoothing, z-loss, softcap, token
scaling,
FP32 gradient accumulation, token accuracy, predicted tokens,
non-contiguous
inputs, partial/frozen-gradient inputs, non-unit upstream gradients,
first-backward ownership handoff, retained-graph repeated backward,
exact-SM100
dispatch, explicit unsupported-configuration errors, custom callback
validation,
non-current-device guards, and public backend dispatch.
Validation environment:
- NVIDIA B200, SM100, driver `580.105.08`
- PyTorch `2.13.0.1+cu130`
- Triton `3.7.1`
- `nvidia-cutlass-dsl` `4.6.0`
- `apache-tvm-ffi` `0.1.12`
## B200 benchmarks
The repository benchmark driver runs the standard FLCE speed and memory
sweeps
for Triton and CuTe DSL in isolated subprocesses:
```bash
python benchmark/scripts/run_cutedsl_compare.py \
--kernel fused_linear_cross_entropy --model llama_3_8b
```
Full sweep results were collected from implementation commit `30eaa4d`.
Scope-cleanup commit `117b286` removes unused GEMM branches without
changing the
executed reference-shape path. The native-only follow-up removes
frontend
fallback ownership from this implementation; its reference-shape smoke
measured
**18.01 ms** and **2,415 MiB**. Latency is the CUDA-event median of ten
full
forward + backward iterations after three warmups (`grad_output=1.0`).
Peak
memory is a warmed isolated-run
`torch.cuda.max_memory_allocated` measurement in MiB, including inputs,
weights,
saved state, and gradients.
`speedup = Triton / CuTe DSL` and `memory ratio = Triton / CuTe DSL`;
values
greater than `1.00x` favor CuTe DSL.
The sweep tables also report CuTe DSL's **effective GEMM throughput**. A
full
forward + backward performs three dense matrix multiplications—logits,
input
gradient, and weight gradient—so the counted work is:
```text
effective GEMM TFLOP/s = 6 * BT * H * V / latency_seconds / 1e12
```
`% peak` uses the B200's 2,250 TFLOP/s dense BF16 tensor-core peak. This
is an
algorithmic end-to-end rate, not an NCU tensor-pipe counter: the
denominator
also includes cross entropy, reductions, chunk scheduling, and
kernel-launch
overhead, so it is a conservative measure of how closely the complete
operator
approaches the GEMM roofline.
### BF16 vocabulary sweep (`BT=8192`, `H=4096`)
| V | CuTe DSL (ms) | Effective GEMM TFLOP/s | % peak | Triton (ms) |
Speedup | CuTe DSL (MiB) | Triton (MiB) | Memory ratio |
|---:|---:|---:|---:|---:|---:|---:|---:|---:|
| 32,000 | 4.52 | 1,425 | 63.3% | 10.70 | 2.37x | 723 | 1,473 | 2.04x |
| 102,400 | 14.19 | 1,453 | 64.6% | 53.90 | 3.80x | 1,960 | 4,260 |
2.17x |
| 128,256 | 17.51 | 1,475 | 65.5% | 121.78 | 6.96x | 2,415 | 5,233 |
2.17x |
| 152,064 | 20.92 | 1,463 | 65.0% | 144.43 | 6.91x | 2,834 | 6,174 |
2.18x |
| 201,088 | 27.71 | 1,461 | 64.9% | 190.27 | 6.87x | 3,697 | 8,116 |
2.20x |
| 262,144 | 36.30 | 1,454 | 64.6% | 484.44 | 13.35x | 4,768 | 10,464 |
2.19x |
### BF16 token sweep (`V=128256`, `H=4096`)
| BT | CuTe DSL (ms) | Effective GEMM TFLOP/s | % peak | Triton (ms) |
Speedup | CuTe DSL (MiB) | Triton (MiB) | Memory ratio |
|---:|---:|---:|---:|---:|---:|---:|---:|---:|
| 1,024 | 2.52 | 1,281 | 56.9% | 117.10 | 46.45x | 2,178 | 5,066 | 2.33x
|
| 2,048 | 4.85 | 1,331 | 59.2% | 117.83 | 24.29x | 2,194 | 5,090 | 2.32x
|
| 4,096 | 9.04 | 1,428 | 63.5% | 119.57 | 13.22x | 2,351 | 5,138 | 2.19x
|
| 8,192 | 17.51 | 1,475 | 65.5% | 121.78 | 6.96x | 2,415 | 5,233 | 2.17x
|
| 16,384 | 35.87 | 1,440 | 64.0% | 133.85 | 3.73x | 2,543 | 5,424 |
2.13x |
| 32,768 | 72.25 | 1,430 | 63.5% | 163.05 | 2.26x | 2,799 | 5,805 |
2.07x |
| 65,536 | 144.83 | 1,426 | 63.4% | 219.88 | 1.52x | 3,312 | 6,569 |
1.98x |
The native path is faster at every BF16 sweep point. The advantage is
largest
where Triton's chunking produces small GEMMs at large vocabulary and
smaller
token counts; it narrows as both implementations approach the large-BT
GEMM
roofline. Effective throughput is stable at 63–66% of dense BF16 peak
across
the vocabulary sweep and once the token sweep reaches `BT >= 4096`.
### FP16 native path
| Dtype | BT | V | CuTe DSL (ms) | Triton (ms) | Speedup | CuTe DSL
(MiB) | Triton (MiB) | Memory ratio |
|:---|---:|---:|---:|---:|---:|---:|---:|---:|
| FP16 | 8,192 | 32,000 | 4.38 | 10.41 | 2.38x | 723 | 1,473 | 2.04x |
| FP16 | 8,192 | 128,256 | 17.92 | 117.51 | 6.56x | 2,415 | 5,233 |
2.17x |
| FP16 | 65,536 | 128,256 | 142.91 | 214.84 | 1.50x | 3,312 | 6,569 |
1.98x |
FP16 uses the same native SM100 path as BF16.
### BF16 features (`BT=8192`, `H=4096`, `V=128256`)
| Configuration | CuTe DSL (ms) | Triton (ms) | Speedup | CuTe DSL (MiB)
| Triton (MiB) | Memory ratio |
|:---|---:|---:|---:|---:|---:|---:|
| Core | 17.98 | 121.79 | 6.77x | 2,415 | 5,233 | 2.17x |
| Bias | 20.56 | 124.77 | 6.07x | 2,604 | 5,233 | 2.01x |
| Label smoothing `0.1` | 17.97 | 121.97 | 6.79x | 2,415 | 5,233 | 2.17x
|
| Z-loss `1e-4` | 17.98 | 121.72 | 6.77x | 2,415 | 5,233 | 2.17x |
| Class weights | 18.02 | 122.40 | 6.79x | 2,416 | 5,233 | 2.17x |
| Token scaling | 32.69 | 132.62 | 4.06x | 2,479 | 5,483 | 2.21x |
| FP32 gradient accumulation | 19.32 | 40.27 | 2.08x | 4,419 | 4,231 |
0.96x |
| Softcap `30` | 22.09 | 125.58 | 5.69x | 2,415 | 5,233 | 2.17x |
All listed features use the same native compute-gradients-in-forward
path.
FP32 gradient accumulation remains faster but uses about 4% more peak
memory
than Triton at this shape.
## Comparison with Dao-AILab Quack FLCE
For peer context, the tables below compare this PR's CuTe DSL sweep with
a
fresh local run of Dao-AILab
[`quack-kernels`](https://github.com/Dao-AILab/quack) **v0.6.2**
(`23c326f`). Quack calls
`quack.linear_cross_entropy.chunked_linear_cross_entropy` with
`chunk_size=4096` and `tuned=True`; these are results from the upstream
Quack
source, not a local reimplementation. No Quack dependency or benchmark
code is
added by this PR.
Both sets use physical GPU 3 (NVIDIA B200), BF16, `H=4096`, three
warmups, ten
measured full forward + backward iterations, and the same
median/peak-memory
methodology. Quack v0.6.2's scaled-exp path is SM90-only, so its public
API
selects the current chunked path on SM100. The exact v0.6.2 source was
run with
CuTe DSL 4.6.0 because its declared 4.6.1 wheel was unavailable from
this
host's package index.
`Quack speedup = CuTe DSL / Quack`, so values above `1.00x` favor Quack.
`Memory ratio = Quack / CuTe DSL`, so values above `1.00x` favor CuTe
DSL.
### Vocabulary sweep (`BT=8192`, `H=4096`)
| V | CuTe DSL (ms) | Quack (ms) | Quack speedup | CuTe DSL (MiB) |
Quack (MiB) | Memory ratio |
|---:|---:|---:|---:|---:|---:|---:|
| 32,000 | 4.52 | 4.30 | 1.05x | 723 | 1,442 | 2.00x |
| 102,400 | 14.19 | 13.92 | 1.02x | 1,960 | 4,192 | 2.14x |
| 128,256 | 17.51 | 17.26 | 1.01x | 2,415 | 5,202 | 2.15x |
| 152,064 | 20.92 | 20.60 | 1.02x | 2,834 | 6,132 | 2.16x |
| 201,088 | 27.71 | 27.02 | 1.03x | 3,697 | 8,050 | 2.18x |
| 262,144 | 36.30 | 35.95 | 1.01x | 4,768 | 10,432 | 2.19x |
### Token sweep (`V=128256`, `H=4096`)
| BT | CuTe DSL (ms) | Quack (ms) | Quack speedup | CuTe DSL (MiB) |
Quack (MiB) | Memory ratio |
|---:|---:|---:|---:|---:|---:|---:|
| 1,024 | 2.52 | 2.12 | 1.19x | 2,178 | 2,335 | 1.07x |
| 2,048 | 4.85 | 3.99 | 1.22x | 2,194 | 2,602 | 1.19x |
| 4,096 | 9.04 | 8.33 | 1.09x | 2,351 | 3,134 | 1.33x |
| 8,192 | 17.51 | 17.26 | 1.01x | 2,415 | 5,202 | 2.15x |
| 16,384 | 35.87 | 34.13 | 1.05x | 2,543 | 5,330 | 2.10x |
| 32,768 | 72.25 | 69.70 | 1.04x | 2,799 | 5,586 | 2.00x |
| 65,536 | 144.83 | 139.96 | 1.03x | 3,312 | 6,099 | 1.84x |
Latest Quack is **1.01-1.05x faster** across the vocabulary sweep and
**1.01-1.22x faster** across the token sweep. CuTe DSL uses **2.00-2.19x
less
peak memory** across vocabulary sizes and **1.07-2.15x less** across
token
counts. At the Llama 3 reference point (`BT=8192`, `V=128256`), latency
is
effectively tied (Quack is **1.01x faster**) while CuTe DSL uses **2.15x
less
peak memory**.
---------
Co-authored-by: Justin Hu <[email protected]>
Co-authored-by: Copilot <[email protected]>
|
||
|
|
7401db04e3 |
[cuTile] Add cutile backend for geglu, layer_norm, fused_linear_jsd, cross_entropy (#1250)
## Summary Adds **cuTile backend implementations** for four Liger-Kernel operators, mirroring the existing Triton kernels under src/liger_kernel/ops/cutile/ops/: - cross_entropy — fused forward + in-place gradient, online numerically-stable logsumexp - fused_linear_jsd — chunked fused linear + JSD (reuses the cuTile jsd_kernel) - geglu — GEGLU (tanh-approx GELU) forward/backward - layer_norm — forward (single-pass mean/var via fold trick) + combined DX/DW/DB backward Also adds a shared element_mul_kernel to cutile/ops/utils.py and wires all new symbols into cutile/ops/__init__.py. ## What's included - cutile/ops/cross_entropy.py — LigerCrossEntropyFunction, cross_entropy_forward/backward - Supports label_smoothing, lse_square_scale (z-loss), per-class weight, softcap, ignore_index, reduction - Adds optional return_token_accuracy / return_predicted_tokens (cross-chunk argmax tracking) — **new logic not present in upstream Triton** - cutile/ops/fused_linear_jsd.py — LigerFusedLinearJSDFunction, fused_linear_jsd_forward/backward - cutile/ops/geglu.py — LigerGELUMulFunction, geglu_forward/backward (aligned/check_bounds=False fast path) - cutile/ops/layer_norm.py — LigerLayerNormFunction, layer_norm_forward/backward (fast path skips DW/DB when not needed) - cutile/ops/utils.py — element_mul_kernel - Benchmark data refresh in benchmark/data/all_benchmark_data.csv and all_benchmark_data_cutile.csv ## Implementation notes - All kernels are faithful ports of the upstream Liger Triton implementations, including the math (online logsumexp, label-smoothing, softcap chain rule sech² = 1 − tanh², GEGLU tanh-approx fwd/bwd, layer-norm c1/c2 reduction) and dtype conventions (loss_1d/Mean/RSTD follow upstream). - Aligned-shape fast paths (check_bounds=False) are used when n_cols is a power of 2 to enable the hardware TMA path. - layer_norm backward uses a persistent grid (grid = sm_count) with per-block DW/DB partials reduced on host. ## Benchmark results <img width="718" height="322" alt="image" src="https://github.com/user-attachments/assets/db40a765-5f81-49e8-aafa-f418782043f6" /> <img width="650" height="401" alt="image" src="https://github.com/user-attachments/assets/e44190ce-3651-4830-a7d7-9538c906f782" /> <img width="649" height="409" alt="image" src="https://github.com/user-attachments/assets/27f98c48-3cc8-4689-986b-1cb551ce63dd" /> <img width="654" height="418" alt="image" src="https://github.com/user-attachments/assets/f39a557c-440b-4812-ac03-1a4cd3715a1e" /> <img width="654" height="402" alt="image" src="https://github.com/user-attachments/assets/88afd7e0-d681-4676-b617-1cf66764114f" /> <img width="666" height="405" alt="image" src="https://github.com/user-attachments/assets/3072d59c-cd1d-48c1-8e68-49b5f2679c7a" /> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- <img width="2198" height="528" alt="image" src="https://github.com/user-attachments/assets/bb6e1a05-8475-4a06-b5c3-f9ebf5be43ef" /> Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: jinmanx (generated by with_the_same_user script) <[email protected]> Co-authored-by: Claude Sonnet 4.6 <[email protected]> Co-authored-by: jinmanx (generated by with_the_same_user script) <[email protected]> Co-authored-by: jinmanx (generated by with_the_same_user script) <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
|
8ef203a42d |
Add cutile jsd (#1228)
## Summary Add optional cuTile backend support for Liger Kernel. This change introduces `LIGER_KERNEL_BACKEND=cutile` backend selection, CUDA-only validation, cuTile JSD operator replacement, optional `cutile` package extras, benchmark labeling as `liger-cutile`, and documentation for enabling cuTile. Related to https://github.com/linkedin/Liger-Kernel/issues/1205 ## Details cuTile is only enabled explicitly through `LIGER_KERNEL_BACKEND=cutile`. Non-CUDA devices raise an error, and missing cuTile dependencies raise an import error with installation instructions. ## Testing Done - Ran JSD correctness tests with cuTile enabled: `LIGER_KERNEL_BACKEND=cutile pytest -s test/transformers/test_jsd.py -v` - Added backend selection test: `test/transformers/test_cutile_backend.py` - Ran targeted ruff checks on modified files. - Ran JSD benchmarks comparing `liger` vs `liger-cutile` on B300, will add B200 in the later commit - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence <img width="1239" height="354" alt="image" src="https://github.com/user-attachments/assets/3abeacdc-dd48-45e3-967f-09f04d254f12" /> --------- Co-authored-by: Vaibhav Jindal <[email protected]> Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]> |
||
|
lowdy1
|
ace042ec55 |
[Benchmark] Refactor benchmark for linear kernels (#1223)
## Summary This PR is a descendant of [[PR #1195](https://github.com/linkedin/Liger-Kernel/pull/1195)] and continues the benchmark refactor effort by migrating linear-space-complexity benchmark scripts to the new reusable `setup` helper infrastructure. ### Benefits * Significantly reduces benchmark script complexity * Makes new benchmark additions easier and more consistent * Centralizes probing and sweep logic * Improves maintainability of benchmark infrastructure ### Notes This PR mainly focuses on kernels with approximately linear memory complexity and migrates them to the new benchmark helper APIs introduced in the parent refactor. - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
lowdy1
|
899233a6af |
[Benckmark] Benchmark refactor example (#1199)
## Summary
The current benchmark scripts contain significant boilerplate when
constructing `common_configs` for `run_benchmarks`. Although
`compute_model_config_sweep_config` and `compute_seq_len_sweep_config`
provide the core sweep logic, each script still:
* Reimplements probe logic
* Manually builds `extra_benchmark_config`
* Duplicates `common_configs` assembly
* Defines redundant helpers like `_resolve_*` and `*_model_config`
variants
This PR removes that duplication by introducing higher-level builders
that standardize how benchmarks are defined.
---
## Proposal
### 1. Introduce higher-level sweep builders
Add two unified helper functions in `benchmark_model_configs.py`:
```python
build_model_config_sweep(...)
build_token_length_sweep(...)
```
These functions:
* Wrap existing sweep utilities (`compute_*_sweep_config`)
* Internally handle memory probing via `setup_fn + forward_fn`
* Automatically construct `extra_benchmark_config` from:
* `model_keys` (dynamic model attributes)
* `extra_configs` (static overrides)
* Return a fully-formed `common_configs` dict
So benchmark scripts reduce to:
```python
common_configs = build_*(...)
run_benchmarks(**common_configs)
```
---
### 2. Standardize kernel definition via `setup_fn`
Instead of manually writing `probe_fn`, all kernels now define:
```python
setup_fn: SingleBenchmarkRunInput -> Tuple[Any, ...]
forward_fn: Tuple[Any, ...] -> torch.Tensor (optional)
```
The builders handle:
```python
setup_out = setup_fn(input)
output = forward_fn(*setup_out)
```
A default is provided:
```python
forward_fn = lambda x, layer: layer(x)
```
This removes duplicated forward/probe logic across scripts.
---
### 3. Eliminate redundant helpers
The following patterns are removed across benchmark scripts:
* `probe_fn` definitions
* `extra_config_fn`
* `_resolve_model_config_*`
* `bench_*_model_config` variants
All are now handled centrally by the builders.
---
## New APIs
### `build_model_config_sweep`
* Sweeps across model configurations (x-axis = model name)
* Keeps total tokens (`B * T`) approximately constant
* Uses `setup_fn + forward_fn` to estimate memory per model
```python
build_model_config_sweep(
kernel_name,
all_model_configs,
setup_fn,
model_keys,
forward_fn=...,
probe_provider="torch",
extra_configs=None,
bt=2048,
overwrite=False,
)
```
---
### `build_token_length_sweep`
* Sweeps across sequence length (x-axis = T)
* Automatically adjusts batch size based on memory estimation
* Uses the same `setup_fn + forward_fn` abstraction
```python
build_token_length_sweep(
kernel_name,
probe_seq_len,
model,
setup_fn,
model_keys,
extra_configs=None,
forward_fn=...,
probe_provider="torch",
x_values_fn=...,
overwrite=False,
)
```
---
## Example (after refactor)
```python
common_configs = build_token_length_sweep(
kernel_name="layer_norm",
probe_seq_len=1024,
model=model,
setup_fn=_setup_layer_norm,
model_keys=["hidden_size", "dtype"],
extra_configs={"eps": 1e-6},
probe_provider="huggingface",
)
common_configs["kernel_providers"] = ["liger", "huggingface"]
run_benchmarks(..., **common_configs)
```
---
## **Benchmark Command Examples**
```bash
python ./benchmark/scripts/benchmark_swiglu.py --sweep-mode model_config [--model llama_3_8b]
python ./benchmark/scripts/benchmark_swiglu.py [--sweep-mode token_length] [--bt 2048]
```
---
### Notes
* `model_config`: sweeps across different model configurations (fixed
total tokens)
* `token_length`: sweeps across sequence lengths / batch sizes (fixed
model) and this is the default and can be omitted
---
- Hardware Type: A100-80G-PCIe
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
|
||
|
Chenguang LiandTcc0403
|
d991472e1e |
[Benchmark]: Add --sweep-mode and --bt to benchmark CLI. (#1163)
This PR follows [PR1162](https://github.com/linkedin/Liger-Kernel/pull/1116) and implements Phase 2. - benchmark_model_configs: replace hidden-size sweep with compute_model_config_sweep_config / ModelConfigSweepConfig; probe each registry model to pick safe batch_size and seq_len for discrete sweeps. - benchmark_geglu / benchmark_swiglu: support model_config sweep across MODEL_REGISTRY via _resolve_model_config_* helpers. - benchmark_dyt: default path sweeps B*T with fixed model dimensions (compute_seq_len_sweep_config); optional model_config sweep; setup uses cfg hidden_size and input.x as BT. - utils: allow string x / x_values for model name indices; extend types. - benchmarks_visualizer: forward extra_config_filter to plotting. - BENCHMARK_GUIDELINES: document D1/D2 sweep patterns and model_config flow. Hardware Type: Atlas 800I A2 - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: Tcc0403 <[email protected]> |
||
|
Chenguang Li
|
c2f4c02015 |
[Test]: Refactor benchmark_geglu with standardized model configs (#1116)
Introduce `benchmark_model_configs.py` with canonical LLM architecture profiles (`LLAMA_2_7B`, `LLAMA_3_8B`) and two utilities: - **`estimate_kernel_bytes_per_token`**: runtime probe that measures peak GPU memory per token for a given kernel, providing a safe upper bound. - **`compute_benchmark_shape`**: derives batch size and sequence length from available device memory, the model profile, and the measured bytes-per-token, replacing hardcoded shapes. Add **`run_speed_benchmark`** and **`run_memory_benchmark`** helpers to `utils.py` to eliminate repeated `triton.testing` / `_test_memory` blocks across scripts. Also add a `--model` CLI argument to `parse_benchmark_script_args` so callers can select a model profile at runtime. Refactor `benchmark_dyt.py`, `benchmark_geglu.py`, and `benchmark_swiglu.py` to: - Extract a **setup factory function** to avoid duplicating layer/tensor construction between speed and memory benchmarks. - Delegate to `run_speed_benchmark` / `run_memory_benchmark`. - Compute `x_values`, `bsz`, and related fields dynamically via `compute_benchmark_shape` instead of hardcoded constants. - Expand memory benchmark modes from `["full"]` to `["full", "forward", "backward"]` for consistency with speed benchmarks. Hardware Type: NVIDIA A100-SXM4-80GB - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
|
b6ed7357bd |
Automate benchmarking - downloading data from GPU to local (#756)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: Manan Shah <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
Andre SlavescuandTcc0403
|
5d25e46827 |
Multi-token Attention (#689)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> Paper: https://arxiv.org/pdf/2504.00927 issue: https://github.com/linkedin/Liger-Kernel/issues/688 <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> tested fp32 + bf16 <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: 3090 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Tcc0403 <[email protected]> |
||
|
Shivam Sahni
|
77ff1a9b7d | Move the checkstyle to [Ruff](https://docs.astral.sh/ruff/) (#483) | ||
|
Golam RabbaniandShao Tang
|
7e3683e23f |
Xpu support (#407)
## Summary Replica of #396 Adds xpu support so all tests, benchmarks etc. run on XPUs or Intel GPUs. ## Details infer_device() function is moved to a separate file and in any file where previously "cuda" was needed, infer_device is imported and "cuda" is replaced with return value of a call to infer_device() ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> A100 80GB PCIe, RTX 3060, Intel Data Center GPU Max 1550 <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Shao Tang <[email protected]> |
||
|
Steven Shimizu
|
b5672a153b |
Refactored benchmark tests (#196)
## Summary
- Refactored kernel benchmarking tests to save data to a single CSV
(`benchmark/data/all_benchmark_data.csv`) that contains more complete
information on how the benchmarking test was setup:
```
kernel_name (e.g. swiglu)
kernel_provider (e.g. liger or huggingface)
kernel_operation_mode (e.g. full, forward, backward)
metric_name (e.g. speed, memory)
metric_unit (e.g. MB)
x_name
x_label (e.g. hidden size)
x_value
y_value_50 (median)
y_value_20 (20th percentile)
y_value_80 (80th percnetile)
extra_benchmark_config_str (e.g. {""B"": 32, ""T"": 512, ""D"": 768, ""dtype"": ""torch.float32""})
gpu_name (e.g. NVIDIA A100-SXM4-80GB)
timestamp
liger_version (e.g. 0.2.1)
```
- Removed existing image files to not clutter up the git history per
previous discussion (https://github.com/linkedin/Liger-Kernel/pull/54)
- Added sample jupyter notebook that can be used to filter and plot the
benchmark data
- Added makefile command to run all benchmarks: `make run-benchmarks`
- Don't overwrite existing data by default if same benchmark. Same
benchmark means same kernel_name, kernel_provider,
kernel_operation_mode, metric_name, x_label, gpu_name, and
extra_benchmark_config_str
## Testing Done
- Ran all benchmarks and spot-checked values with existing benchmark
data
- Hardware Type: A100
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
|