mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Arup DeandCopilot
96233e3d9b
# feat(cutile): cuTile backends for RMSNorm, LayerNorm, Softmax, JSD,
FusedLinearJSD
**Head:** `arde171:arde/liger-E-cutile` → **Base:** `main`. Final PR of
the multi-backend series; **#1419 (CuTeDSL CE) is now merged**, so this
diff is **cuTile-only** — no rebase caveat anymore.
## What this adds
Registers a **cuTile** backend family for five ops, gated strictly to
Blackwell. All impls declare **`min_cc=(10, 0)`** → only reachable on
sm_100+ (B200/B300); Hopper/Ampere transparently keep using
Triton/CuTeDSL.
| file | + |
|---|---|
| `ops/backends/_cutile/rms_norm.py` | 1221 |
| `ops/backends/_cutile/layer_norm.py` | 626 |
| `ops/backends/_cutile/softmax.py` | 448 |
| `ops/backends/_cutile/jsd.py` | 319 |
| `ops/backends/_cutile/fused_linear_jsd.py` | 64 |
| `ops/backends/_cutile/__init__.py` | 11 |
| `functional.py` | 6 (append `_cutile.*` to 5 op declarations) |
| `ops/fused_linear_jsd.py` | 19 (`_max_capability(device=None)`) |
| `testing/_op_helpers.py` | 5 (`strict_grad` flag) |
| `test/backends/test_cutile_select_mode.py` | 83 |
| `test/cutile/test_cutile_backends_parity.py` | 799 (113-case parity
matrix) |
| `test/ops/{test_rms_norm,test_layer_norm}.py` | 44 |
No CuTeDSL file touched; purely additive registration behind a Blackwell
gate.
## Default-selection is benchmark-justified (nothing defaults to a
slower kernel)
Ranks are chosen so cuTile is **only ever opt-in** — it is never
auto-selected over a faster existing path. Measured median latency
(bf16, fwd+bwd) on **B200 (sm_100)**:
| op | rank | Triton | cuTile | CuTeDSL | auto-selected |
|---|---|---|---|---|---|
| `rms_norm` | 80 | 0.262 ms | 0.473 ms | **0.190 ms** | CuTeDSL |
| `layer_norm` | 80 | 0.271 ms | 0.623 ms | **0.219 ms** | CuTeDSL |
| `softmax` | 80 | **0.151 ms** | 0.188 ms | 0.156 ms | Triton |
| `jsd` | 25 | 0.802 ms | 0.661 ms | **0.583 ms** | CuTeDSL |
| `fused_linear_jsd` | 25 | — | — | — | CuTeDSL |
- **Norms → `preference_rank=80`** (above Triton's 50): cuTile is
0.44–0.55× of Triton on B200, so it must stay opt-in.
- **JSD / FusedLinearJSD → `preference_rank=25`**: both cuTile (1.21×)
and CuTeDSL (1.38×) beat Triton, but CuTeDSL wins, so CuTeDSL stays the
default and cuTile is opt-in below it.
## Testing (deterministic `-p no:randomly`)
Full `test/backends test/ops test/cutile` + the 113-case
CuTeDSL↔cuTile↔Torch parity matrix:
- **B300 (cc 10.3):** `1149 passed, 48 skipped, 0 failed` + `113/113`
parity.
- **B200 (cc 10.0):** `1149 passed, 48 skipped, 0 failed` + `113/113`
parity.
- **H200 (cc 9.0):** green (cuTile `min_cc=(10,0)` correctly skips on
Hopper).
- ruff check + format clean. Mixed-precision affine (dW/dB in their own
dtypes), device-safe launches (all kernels use the input tensor's
device), and `strict_grad` elementwise dW/dB assertions added.
Co-authored-by: Copilot <[email protected]>
---------
Co-authored-by: Copilot <[email protected]>