mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
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:
@@ -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:
|
||||
|
||||
@@ -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.
|
||||
#
|
||||
|
||||
Reference in New Issue
Block a user