mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Vaibhav JindalandClaude Opus 4.6
39a1f454f3
This code change was generated end-to-end by Claude Code using the `liger-kernel-dev` skill introduced in #1170. ## Summary Adds triton implementation for relu squared activation function. This is used for NVIDIA's nemotron models. (https://github.com/huggingface/transformers/blob/78bdaf0b39c29737b9ca48a274ef4a34bdafd4d1/src/transformers/activations.py#L206) ### Files added/modified | File | Description | |------|-------------| | `src/liger_kernel/ops/relu_squared.py` | **NEW** — Triton forward/backward kernels + `LigerReLUSquaredFunction` autograd Function | | `src/liger_kernel/transformers/relu_squared.py` | **NEW** — `LigerReLUSquared` nn.Module wrapper | | `src/liger_kernel/transformers/functional.py` | **MODIFIED** — added `liger_relu_squared()` functional API | | `src/liger_kernel/ops/__init__.py` | **MODIFIED** — export `LigerReLUSquaredFunction` + forward/backward | | `src/liger_kernel/transformers/__init__.py` | **MODIFIED** — export `LigerReLUSquared` + added to `__all__` | | `test/transformers/test_relu_squared.py` | **NEW** — 18 unit tests (correctness + functional API, float32 + bf16, multiple shapes) | | `benchmark/scripts/benchmark_relu_squared.py` | **NEW** — speed + memory benchmarks | | `benchmark/data/all_benchmark_data.csv` | **MODIFIED** — benchmark results appended | ### Kernel design - **Tier 1 (element-wise)** — simplest kernel tier, one program per row - **Forward**: fused `relu + square` in a single kernel pass - **Backward**: `dX = dY * 2 * relu(X)`, reads saved X without modifying it - **No float32 upcasting needed** — relu and square are numerically stable - Uses `calculate_settings(n_cols)` for BLOCK_SIZE/num_warps auto-tuning ## Benchmark Results (NVIDIA H100 80GB HBM3, bfloat16, M=4096) ### Speed | Hidden Size | Mode | Liger (ms) | PyTorch (ms) | Speedup | |-------------|------|-----------|-------------|---------| | 4096 | Forward | 0.028 | 0.053 | **1.9x** | | 4096 | Backward | 0.039 | 0.126 | **3.2x** | | 4096 | Full | 0.087 | 0.178 | **2.0x** | | 16384 | Forward | 0.094 | 0.186 | **2.0x** | | 16384 | Backward | 0.137 | 0.453 | **3.3x** | | 16384 | Full | 0.229 | 0.635 | **2.8x** | ### Memory | Hidden Size | Liger (MB) | PyTorch (MB) | Savings | |-------------|-----------|-------------|---------| | 4096 | 0.039 | 0.063 | **37.5%** | | 16384 | 0.156 | 0.250 | **37.5%** | ### Speed — Forward <!-- INSERT relu_squared_speed_forward.png --> <img width="1000" height="600" alt="relu_squared_speed_forward" src="https://github.com/user-attachments/assets/747e0ede-8d97-449c-9505-a1ba21f5470d" />. ### Speed — Backward <!-- INSERT relu_squared_speed_backward.png --> <img width="1000" height="600" alt="relu_squared_speed_backward" src="https://github.com/user-attachments/assets/033f3373-f1f5-44f0-b630-393b82bfd93c" /> ### Speed — Full (Forward + Backward) <!-- INSERT relu_squared_speed_full.png --> <img width="1000" height="600" alt="relu_squared_speed_full" src="https://github.com/user-attachments/assets/cb32cc32-6a39-4b2a-a61e-2df92c88414e" /> ### Memory — Full <!-- INSERT relu_squared_memory_full.png --> <img width="1000" height="600" alt="relu_squared_memory_full" src="https://github.com/user-attachments/assets/0128bade-ec3b-419a-9885-f7e37ca8c400" /> ## Testing Done Hardware Type: NVIDIA H100 80GB HBM3 - [x] 18/18 unit tests pass on H100 (float32 + bfloat16, 6 shape configs including non-power-of-2) - [x] `ruff check` and `ruff format` pass - [x] Speed benchmarks run successfully (forward, backward, full) - [x] Memory benchmarks run successfully --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>