linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(triton): pin fp64 scalar kernel params to fp32 in swiglu and rms_norm (#1350)

Browse Source 
## Summary

Hand-written Triton kernels get their non-constexpr Python-float params
specialized differently by eager Triton and by Inductor. Eager
specializes them to `fp32`; when the same kernel is launched from inside
`torch.compile`, Inductor specializes them to `fp64`:

```
# TORCH_LOGS=output_code, _swiglu_forward_kernel as generated by Inductor
triton_meta={'signature': {'a_ptr': '*bf16', ..., 'gate_multiplier': 'fp64', ...}}
```

Since `x.to(tl.float32) * gate_multiplier` mixes an fp32 tensor with an
fp64 scalar, Triton promotes the **whole expression to float64** —
including the transcendentals (`tl.sigmoid` in swiglu, `rsqrt` in
rms_norm). These kernels are memory bound, so the promotion roughly
halves their throughput.

Isolated repro at 8192x3072 bf16 on H100:

| scalar typing | time | bandwidth |
|---|---|---|
| fp32 (eager path) | 56.1 us | 2693 GB/s |
| fp64 (Inductor path) | 108.5 us | 1392 GB/s |
| fp64 + `.to(tl.float32)` | **56.0 us** | **2694 GB/s** |

The fix is an explicit `.to(tl.float32)` on the scalar inside the kernel
— a no-op on the eager path (already fp32) and a full recovery on the
compiled path. `rms_norm` already did exactly this for
`_CASTING_MODE_NONE`; this extends it to the LLAMA/GEMMA casting modes
and to the affine `offset` used by both backward kernels.

## Details

This was found while benchmarking Liger against `torch.compile` on
Qwen3-0.6B. Worth noting what it was **not**: there are no graph breaks
in either config (Liger actually lowers op count 1266 -> 844), and
standalone the Liger SwiGLU kernel already matched Inductor's exactly
(both 56us, ~80% of roofline). The regression was `torch.compile`
silently degrading Liger's own kernels via scalar typing.

Per-kernel GPU time, Qwen3-0.6B trunk forward @ 8x1024, H100 bf16:

| kernel | before | after | compile-only equivalent |
|---|---|---|---|
| `_swiglu_forward_kernel` | 3.017 ms | **1.361 ms** | 1.316 ms |
| `_rms_norm_forward_kernel` | 0.848 ms | **0.655 ms** | — |
| total GPU | 24.396 ms | **22.958 ms** | 23.395 ms |

End-to-end speedup vs eager HF baseline (Liger patched with `rope +
rms_norm + swiglu` only; CE / fused-linear-CE off so the LM head is
identical across configs):

| shape | mode | `torch.compile` | `liger+compile` before |
`liger+compile` after |
|---|---|---|---|---|
| 8x1024 | trunk fwd | 1.870x | 1.856x | **2.075x** |
| 8x2048 | trunk fwd | 1.557x | 1.509x | **1.674x** |
| 4x2048 | trunk fwd | 1.578x | 1.534x | **1.727x** |
| 8x1024 | fwd+bwd | 1.978x | 1.845x | **2.008x** |
| 8x2048 | fwd+bwd | 1.651x | 1.554x | **1.691x** |

So stacking Liger with `torch.compile` was previously a regression
relative to `torch.compile` alone, and is now a win. Eager-only Liger is
unchanged (27.765 ms vs 27.755 ms before), confirming the cast is free
on that path.

Numerical drift versus the eager baseline also **improves**, since the
compiled path now does the same fp32 arithmetic eager does:

```
liger        : rel=2.01e-02  argmax_agree=100.0%
compile      : rel=1.92e-02  argmax_agree=100.0%
liger+compile: rel=4.77e-02  ->  2.51e-02   argmax_agree=100.0%
```

Two things I deliberately left out of this PR:

1. **The same pattern likely affects other ops.** An AST scan turned up
~80 float-scalar kernel params across ~20 kernels (`layer_norm`,
`group_norm`, `poly_norm`, `dyt`, `cross_entropy` softcap, `grpo_loss`
temperature, `tvd` scale, ...). I only fixed the two ops I could measure
end-to-end. Happy to do a sweep in a follow-up.
2. **Unrelated:** `torch.compile` with Liger RoPE requires
`dynamic=False` — with dynamic shapes Dynamo hits an internal assertion
in `higher_order_ops.py` (`assert subgraph_vt.is_tensor() or
isinstance(subgraph_vt, SymNodeVariable)`) while tracing
`LigerRopeFunction`, on the *second* distinct input shape. That's a
separate issue, not addressed here.

## Testing Done

Correctness on the touched ops, plus convergence:

- `pytest test/transformers/test_swiglu.py
test/transformers/test_rms_norm.py` -> **113 passed**, 4 skipped, 12
xfailed
- `pytest test/transformers/test_fused_add_rms_norm.py
test/transformers/test_monkey_patch.py` -> **100 passed**, 2 xfailed
- `pytest test/convergence/bf16/test_mini_models.py -k "qwen3 or llama"`
-> **10 passed**
- `make checkstyle` -> passed

- Hardware Type: **H100 80GB HBM3**
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

> Note: I ran the targeted suites listed above rather than the full
`make test` /
> `make test-convergence` sweeps, so I've left those two boxes unchecked
rather than
> claim more coverage than I actually have. Relying on CI for the full
run.

Co-authored-by: Copilot <[email protected]>
This commit is contained in:
Vaibhav Jindal
2026-08-07 18:55:27 +00:00
committed by GitHub
co-authored by Copilot
parent 4c80e19130
commit 11702da962
2 changed files with 30 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/liger_kernel/ops/rms_norm.py
+16 -2
View File
@@ -106,6 +106,14 @@ def _rms_norm_forward_kernel(
if casting_mode == _CASTING_MODE_NONE:
eps = eps.to(X_row_dtype)
offset = offset.to(X_row_dtype)
else:
# Scalar kernel params are specialized to fp32 by eager Triton but to
# fp64 by Inductor when this kernel is launched from inside
# torch.compile. An fp64 scalar silently promotes mean_square, rsqrt
# and the weight multiply to float64, roughly halving throughput.
# Pinning to fp32 is a no-op in eager and keeps both paths identical.
eps = eps.to(tl.float32)
offset = offset.to(tl.float32)
mean_square = tl.sum(X_row * X_row, axis=0) / n_cols
rstd = rsqrt(mean_square + eps)
@@ -171,7 +179,8 @@ def _rms_norm_backward_kernel(
if elementwise_affine:
W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
W_row = W_row + offset
# Pin the fp64 scalar Inductor passes to fp32 (see _rms_norm_forward_kernel).
W_row = W_row + offset.to(tl.float32)
for row_idx in range(row_start, row_end):
dy_base = dY_ptr + row_idx * dY_row_stride
@@ -280,6 +289,10 @@ def _block_rms_norm_forward_kernel(
if casting_mode == _CASTING_MODE_NONE:
eps = eps.to(X_row_dtype)
offset = offset.to(X_row_dtype)
else:
# See _rms_norm_forward_kernel: pin fp64 scalars from Inductor to fp32.
eps = eps.to(tl.float32)
offset = offset.to(tl.float32)
mean_square = tl.sum(X_row * X_row, axis=1) / n_cols
rstd = rsqrt(mean_square + eps)
@@ -348,7 +361,8 @@ def _block_rms_norm_backward_kernel(
dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
W_row = tl.load(W_ptr + col_offsets, mask=col_mask, other=0.0)
W_row = W_row + offset
# Pin the fp64 scalar Inductor passes to fp32 (see _rms_norm_forward_kernel).
W_row = W_row + offset.to(tl.float32)
for start in range(pid * BLOCK_ROW, n_rows, NUM_SMS * BLOCK_ROW):
row_idx = start + tl.arange(0, BLOCK_ROW)
src/liger_kernel/ops/swiglu.py
+14 -6
View File
@@ -36,6 +36,14 @@ def silu(x):
return x * tl.sigmoid(x)
# NOTE on `gate_multiplier.to(tl.float32)` below: scalar kernel params are
# specialized to fp32 by eager Triton but to fp64 by Inductor when these kernels
# are launched from inside torch.compile. An fp64 scalar silently promotes the
# whole silu/sigmoid chain to float64, which halves throughput (108us vs 56us at
# 8192x3072 on H100). The explicit cast is a no-op in eager and keeps the
# compiled path identical.
@triton.jit
def _swiglu_forward_kernel(
a_ptr, b_ptr, c_ptr, stride, gate_multiplier, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
@@ -51,7 +59,7 @@ def _swiglu_forward_kernel(
mask = col_offsets < n_cols
# sigmoid requires type float32
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32) * gate_multiplier
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32) * gate_multiplier.to(tl.float32)
b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
c_row = silu(a_row).cast(b_row.dtype) * b_row
tl.store(c_ptr + col_offsets, c_row, mask=mask)
@@ -73,7 +81,7 @@ def _swiglu_backward_kernel(
dc_row = tl.load(dc_ptr + col_offsets, mask=mask, other=0)
# sigmoid requires type float32
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32) * gate_multiplier
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32) * gate_multiplier.to(tl.float32)
b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
# recomputation to save memory. a_row already holds a * gate_multiplier.
@@ -81,7 +89,7 @@ def _swiglu_backward_kernel(
silu_a = a_row * sig_a
db_row = dc_row * silu_a
# chain rule pulls an extra factor of gate_multiplier through the pre-activation scaling
da_row = dc_row * (silu_a * (1 - sig_a) + sig_a) * b_row * gate_multiplier
da_row = dc_row * (silu_a * (1 - sig_a) + sig_a) * b_row * gate_multiplier.to(tl.float32)
tl.store(a_ptr + col_offsets, da_row, mask=mask)
tl.store(b_ptr + col_offsets, db_row, mask=mask)
@@ -102,7 +110,7 @@ def _swiglu_forward_kernel_tiled(a_ptr, b_ptr, c_ptr, stride, gate_multiplier, n
mask = col_offsets < n_cols
# sigmoid requires type float32
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32) * gate_multiplier
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32) * gate_multiplier.to(tl.float32)
b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
c_row = silu(a_row).cast(b_row.dtype) * b_row
tl.store(c_ptr + col_offsets, c_row, mask=mask)
@@ -124,7 +132,7 @@ def _swiglu_backward_kernel_tiled(dc_ptr, a_ptr, b_ptr, stride, gate_multiplier,
dc_row = tl.load(dc_ptr + col_offsets, mask=mask, other=0)
# sigmoid requires type float32
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32) * gate_multiplier
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32) * gate_multiplier.to(tl.float32)
b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
# recomputation to save memory. a_row already holds a * gate_multiplier.
@@ -132,7 +140,7 @@ def _swiglu_backward_kernel_tiled(dc_ptr, a_ptr, b_ptr, stride, gate_multiplier,
silu_a = a_row * sig_a
db_row = dc_row * silu_a
# chain rule pulls an extra factor of gate_multiplier through the pre-activation scaling
da_row = dc_row * (silu_a * (1 - sig_a) + sig_a) * b_row * gate_multiplier
da_row = dc_row * (silu_a * (1 - sig_a) + sig_a) * b_row * gate_multiplier.to(tl.float32)
tl.store(a_ptr + col_offsets, da_row, mask=mask)
tl.store(b_ptr + col_offsets, db_row, mask=mask)