linkedin / linkedin/Liger-Kernel
[New backend] cuTile backend for Liger kernels — performance data on B200 + integration interest
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 6.6k
- Forks
- 603
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 47
Description
## Background
Hi Liger team :wave:
We've been building **[TileGym](https://github.com/NVIDIA/tilegym)**, helpful kernel tutorials and examples for tile-based GPU programming.
Over the past few weeks, we have implemented cuTile Python versions of most Liger kernels internally and want to release them in TileGym. We wanted to share some early performance numbers and start a conversation about whether there might be an opportunity to integrate a cuTile backend into Liger down the road.
## What we've implemented
We have working internal implementations of the following kernels (tested on **NVIDIA B200, bfloat16, Llama-3-8B config** — hidden_size=4096, intermediate_size=14336):
```
cross_entropy, dyt, fused_add_rms_norm, fused_linear_cross_entropy, fused_linear_jsd, fused_neighborhood_attention, geglu, group_norm, jsd, kl_div, layer_norm, llama4_rope, poly_norm, qwen2vl_mrope, rms_norm, rope, softmax, sparsemax, swiglu, tiled_geglu, tiled_swiglu, tvd
```
We currently have tested cuTile-Python backends for Liger-equivalent kernels internally(tested on **NVIDIA B200, 1650 MHz, bfloat16, Llama-3-8B config**, triton version: 3.7.0+git7d075612). Below are some of our very promising results.
We use the following command to test it.
```
python benchmark_${kernel}.py --model llama_3_8b --sweep-mode token_length
```
```
e.g. jsd
Full (Fwd + Bwd) Speed + Memory
┌───────┬─────────────┬──────────────┬──────────────┬─────────┬───────────┬────────────┬───────────┐
│ B×T │ Torch Speed │ Liger Speed │ CuTile Speed │ Speedup │ Liger Mem │ CuTile Mem │ Mem Ratio │
│ │ │ │ │(vs Liger)│ │ │ │
├───────┼─────────────┼──────────────┼──────────────┼─────────┼───────────┼────────────┼───────────┤
│ 1,024 │ 6.12 ms │ 6.10 ms │ 1.39 ms │ 4.38× │ 3,012 MB │ 3,012 MB │ 1.00× │
│ 2,048 │ 12.24 ms │ 10.56 ms │ 2.57 ms │ 4.11× │ 6,012 MB │ 6,012 MB │ 1.00× │
│ 4,096 │ 24.40 ms │ 21.13 ms │ 4.95 ms │ 4.27× │ 12,024 MB │ 12,024 MB │ 1.00× │
│ 8,192 │ 48.89 ms │ 42.14 ms │ 9.79 ms │ 4.30× │ 24,048 MB │ 24,048 MB │ 1.00× │
└───────┴─────────────┴──────────────┴──────────────┴─────────┴───────────┴────────────┴───────────┘
Consistent 4× speedup; memory bit-for-bit identical. Here we use a different BLOCK_SIZE for cuTile
```
## Proposed Integration Plan
**Step 1 — TileGym side** (we own this entirely)
Release the CuTile implementations under tilegym/suites/liger/, publish to PyPI as part of tilegym. No changes needed in Liger for this step.
**Step 2 — Liger side** (we'd like to contribute a PR, pending your approval)
Add an optional CuTile backend that activates when tilegym is installed:
```
liger_kernel/ops/backends/cutile/
__init__.py # try: import tilegym; except: pass
group_norm.py # thin dispatch wrappers
jsd.py
...
```
## Questions for the team
1. Is the pip install tilegym + optional dispatch model acceptable to you? Or would you prefer a different integration shape?
2. What are your requirements for adding an optional dependency (testing policy, supported platforms)?
3. Happy to start small — would you be open to a draft PR with just jsd as a proof of concept?
Thanks for all the great work on Liger
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reviewing the proposed optional backend layout under liger_kernel/ops/backends/cutile/ and the jsd proof-of-concept scope. Use the listed benchmark command and B200 results as context, then resolve the questions about optional tilegym installation, testing policy, and supported platforms. Done means the team has agreed on an integration shape and contribution requirements.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100