[New backend] cuTile backend for Liger kernels — performance data on B200 + integration interest

[xjmxyt]

Background

Hi Liger team 👋

We've been building 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

[vaibhavjindal]

Hi @xjmxyt , thanks for opening up this issue. The numbers on jsd kernel look very promising and we would love to see how we can work on this integration and understand how cuTile compares with triton.

Reply to questions:

1. We need to discuss this further on our end to see what could be a good integration idea. But for now, we can start with a draft PR and see if the integration is working as expected. If I understand correctly, you are proposing to use the tilegym kernels when the gpu is blackwell and tilegym is installed, right?

2. We don't have a very strict testing policy for this. We can work on this as and when we finalize on the integration.

3. Yes, let's start with a draft PR for jsd kernel for an initial PoC. Our team can review it and then discuss on the next steps after this.