fix(mlp): keep liger_kernel.ops importable without TensorDescriptor (#1421)

## Summary

- `LigerMLP` (#1357) does a hard `from triton.tools.tensor_descriptor
import TensorDescriptor`. That module first shipped in official Triton
**3.4.0**, so any earlier / vendor build fails while importing
`liger_kernel.ops`.
- `ops/__init__.py` always loads default kernels before backend overlay,
so one missing import takes down every op (RMSNorm, Softmax, RoPE, …),
not just MLP.
- Make the import optional. `LigerMLPFunction.forward` raises a clear
`RuntimeError` if the API is missing. GPU / Triton >= 3.4 is unchanged.
- Skip `test/transformers/test_mlp.py` when
`triton.tools.tensor_descriptor` is not installed.

## Testing Done

- `pytest test/transformers/ --collect-only` on a Triton 3.2 build:
collection succeeds (no longer 33 collection errors).
- `test_mlp.py` is skipped when the API is absent; other transformers
tests collect and run as before.
This commit is contained in:
Rain
2026-08-30 13:02:20 +00:00
committed by GitHub
parent 00e72b86ca
commit 777799588a
2 changed files with 11 additions and 1 deletions
+9 -1
View File
@@ -22,7 +22,11 @@ import torch.nn.functional as F
import triton
import triton.language as tl
from triton.tools.tensor_descriptor import TensorDescriptor
# Optional: not all Triton builds ship triton.tools.tensor_descriptor.
try:
from triton.tools.tensor_descriptor import TensorDescriptor
except ModuleNotFoundError:
TensorDescriptor = None
@triton.jit
@@ -631,6 +635,10 @@ class LigerMLPFunction(torch.autograd.Function):
gate_multiplier=1.0,
down_multiplier=1.0,
):
if TensorDescriptor is None:
raise RuntimeError(
"LigerMLP requires triton.tools.tensor_descriptor, which is not available in this Triton build."
)
assert input.is_cuda and gate_weight.is_cuda and up_weight.is_cuda and down_weight.is_cuda
input, gate_weight, up_weight, down_weight = _check_inputs(input, gate_weight, up_weight, down_weight)
# Note:
+2
View File
@@ -19,6 +19,8 @@ except ImportError:
device = infer_device()
pytest.importorskip("triton.tools.tensor_descriptor", reason="LigerMLP requires triton.tools.tensor_descriptor")
DIFF_THRESHOLD = 1e-5
# Why the FalconH1 tests use a looser tolerance only for bf16.
#