Update import path for ops in liger_kernel/transformers/*

Signed-off-by: Tcc0403 <[email protected]>

src/liger_kernel/ops/__init__.py

@@ -1,6 +1,5 @@
 import importlib
 import inspect
-import logging
 
 from liger_kernel.ops.backends import VENDOR_REGISTRY
 from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
@@ -28,7 +27,6 @@ from liger_kernel.ops.tiled_mlp import LigerTiledMLPFunction
 from liger_kernel.ops.tvd import LigerTVDLossFunction
 from liger_kernel.utils import infer_device
 
-logger = logging.getLogger(__name__)
 device = infer_device()
 
 __all__ = [
@@ -59,7 +57,6 @@ __all__ = [
 _globals = globals()
 
 if device in VENDOR_REGISTRY:
-    logger.debug(f"Loading Liger Kernel ops for device: {device}")
     backend_module_path = VENDOR_REGISTRY[device].module_path
     module = importlib.import_module(backend_module_path)
     ops = []

src/liger_kernel/transformers/cross_entropy.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import torch
 
-from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
+from liger_kernel.ops import LigerCrossEntropyFunction
 from liger_kernel.transformers.functional import CrossEntropyOutput
 
 

src/liger_kernel/transformers/dyt.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.dyt import LigerDyTFunction
+from liger_kernel.ops import LigerDyTFunction
 
 
 class LigerDyT(nn.Module):

src/liger_kernel/transformers/functional.py

@@ -3,26 +3,26 @@ from typing import Optional
 
 import torch
 
-from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
-from liger_kernel.ops.dyt import LigerDyTFunction
-from liger_kernel.ops.fused_add_rms_norm import LigerFusedAddRMSNormFunction
-from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
-from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
-from liger_kernel.ops.fused_neighborhood_attention import LigerFusedNeighborhoodAttentionFunction
-from liger_kernel.ops.geglu import LigerGELUMulFunction
-from liger_kernel.ops.group_norm import LigerGroupNormFunction
-from liger_kernel.ops.jsd import LigerJSDFunction
-from liger_kernel.ops.kl_div import LigerKLDivLossFunction
-from liger_kernel.ops.layer_norm import LigerLayerNormFunction
-from liger_kernel.ops.multi_token_attention import LigerMultiTokenAttentionFunction
-from liger_kernel.ops.poly_norm import LigerPolyNormFunction
-from liger_kernel.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction
-from liger_kernel.ops.rms_norm import LigerRMSNormFunction
-from liger_kernel.ops.rope import LigerRopeFunction
-from liger_kernel.ops.softmax import LigerSoftmaxFunction
-from liger_kernel.ops.sparsemax import LigerSparsemaxFunction
-from liger_kernel.ops.swiglu import LigerSiLUMulFunction
-from liger_kernel.ops.tvd import LigerTVDLossFunction
+from liger_kernel.ops import LigerCrossEntropyFunction
+from liger_kernel.ops import LigerDyTFunction
+from liger_kernel.ops import LigerFusedAddRMSNormFunction
+from liger_kernel.ops import LigerFusedLinearCrossEntropyFunction
+from liger_kernel.ops import LigerFusedLinearJSDFunction
+from liger_kernel.ops import LigerFusedNeighborhoodAttentionFunction
+from liger_kernel.ops import LigerGELUMulFunction
+from liger_kernel.ops import LigerGroupNormFunction
+from liger_kernel.ops import LigerJSDFunction
+from liger_kernel.ops import LigerKLDivLossFunction
+from liger_kernel.ops import LigerLayerNormFunction
+from liger_kernel.ops import LigerMultiTokenAttentionFunction
+from liger_kernel.ops import LigerPolyNormFunction
+from liger_kernel.ops import LigerQwen2VLMRopeFunction
+from liger_kernel.ops import LigerRMSNormFunction
+from liger_kernel.ops import LigerRopeFunction
+from liger_kernel.ops import LigerSiLUMulFunction
+from liger_kernel.ops import LigerSoftmaxFunction
+from liger_kernel.ops import LigerSparsemaxFunction
+from liger_kernel.ops import LigerTVDLossFunction
 
 
 @dataclass

src/liger_kernel/transformers/fused_add_rms_norm.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.fused_add_rms_norm import LigerFusedAddRMSNormFunction
+from liger_kernel.ops import LigerFusedAddRMSNormFunction
 
 
 class LigerFusedAddRMSNorm(nn.Module):

src/liger_kernel/transformers/fused_linear_cross_entropy.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import torch
 
-from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
+from liger_kernel.ops import LigerFusedLinearCrossEntropyFunction
 from liger_kernel.transformers.functional import CrossEntropyOutput
 
 

src/liger_kernel/transformers/fused_linear_jsd.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import torch
 
-from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
+from liger_kernel.ops import LigerFusedLinearJSDFunction
 
 
 class LigerFusedLinearJSD(torch.nn.Module):

src/liger_kernel/transformers/fused_neighborhood_attention.py

@@ -5,7 +5,7 @@ from typing import Optional
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.fused_neighborhood_attention import LigerFusedNeighborhoodAttentionFunction
+from liger_kernel.ops import LigerFusedNeighborhoodAttentionFunction
 
 
 class LigerFusedNeighborhoodAttention(nn.Module):

src/liger_kernel/transformers/geglu.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from liger_kernel.ops.geglu import LigerGELUMulFunction
+from liger_kernel.ops import LigerGELUMulFunction
 
 
 class LigerGEGLUMLP(nn.Module):

src/liger_kernel/transformers/group_norm.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.group_norm import LigerGroupNormFunction
+from liger_kernel.ops import LigerGroupNormFunction
 
 
 class LigerGroupNorm(nn.Module):

src/liger_kernel/transformers/grpo_loss.py

@@ -1,7 +1,7 @@
 import torch
 
 from liger_kernel.chunked_loss.fused_linear_ppo import LigerFusedLinearPPOBase
-from liger_kernel.ops.grpo_loss import GrpoLossFunction
+from liger_kernel.ops import GrpoLossFunction
 
 
 def triton_grpo_loss(

src/liger_kernel/transformers/jsd.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import torch
 
-from liger_kernel.ops.jsd import LigerJSDFunction
+from liger_kernel.ops import LigerJSDFunction
 
 
 class LigerJSD(torch.nn.Module):

src/liger_kernel/transformers/kl_div.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from liger_kernel.ops.kl_div import LigerKLDivLossFunction
+from liger_kernel.ops import LigerKLDivLossFunction
 
 
 class LigerKLDIVLoss(nn.KLDivLoss):

src/liger_kernel/transformers/layer_norm.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.layer_norm import LigerLayerNormFunction
+from liger_kernel.ops import LigerLayerNormFunction
 
 
 class LigerLayerNorm(nn.Module):

src/liger_kernel/transformers/llama4_rope.py

@@ -5,7 +5,7 @@ Supports both text and vision RoPE variants with fused operations for optimal pe
 
 import torch
 
-from liger_kernel.ops.llama4_rope import LigerLlama4RopeFunction
+from liger_kernel.ops import LigerLlama4RopeFunction
 
 
 def liger_llama4_text_rotary_pos_emb(

src/liger_kernel/transformers/multi_token_attention.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 
 from torch.nn.modules.utils import _pair
 
-from liger_kernel.ops.multi_token_attention import LigerMultiTokenAttentionFunction
+from liger_kernel.ops import LigerMultiTokenAttentionFunction
 
 
 class LigerMultiTokenAttention(nn.Module):

src/liger_kernel/transformers/poly_norm.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.poly_norm import LigerPolyNormFunction
+from liger_kernel.ops import LigerPolyNormFunction
 
 
 class LigerPolyNorm(nn.Module):

src/liger_kernel/transformers/qwen2vl_mrope.py

@@ -1,4 +1,4 @@
-from liger_kernel.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction
+from liger_kernel.ops import LigerQwen2VLMRopeFunction
 
 
 def liger_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1):

src/liger_kernel/transformers/rms_norm.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.rms_norm import LigerRMSNormFunction
+from liger_kernel.ops import LigerRMSNormFunction
 
 
 class LigerRMSNorm(nn.Module):

src/liger_kernel/transformers/rope.py

@@ -2,7 +2,7 @@ from typing import Tuple
 
 import torch
 
-from liger_kernel.ops.rope import LigerRopeFunction
+from liger_kernel.ops import LigerRopeFunction
 
 
 def liger_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):

src/liger_kernel/transformers/sparsemax.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.sparsemax import LigerSparsemaxFunction
+from liger_kernel.ops import LigerSparsemaxFunction
 
 
 class LigerSparsemax(nn.Module):

src/liger_kernel/transformers/swiglu.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from liger_kernel.ops.swiglu import LigerSiLUMulFunction
+from liger_kernel.ops import LigerSiLUMulFunction
 
 
 class LigerSwiGLUMLP(nn.Module):

src/liger_kernel/transformers/tiled_mlp.py

@@ -2,9 +2,9 @@ from typing import Optional
 
 import torch.nn as nn
 
-from liger_kernel.ops.geglu import LigerGELUMulFunction
-from liger_kernel.ops.swiglu import LigerSiLUMulFunction
-from liger_kernel.ops.tiled_mlp import apply_tiled_mlp
+from liger_kernel.ops import LigerGELUMulFunction
+from liger_kernel.ops import LigerSiLUMulFunction
+from liger_kernel.ops import apply_tiled_mlp
 
 
 class LigerTiledGEGLUMLP(nn.Module):

src/liger_kernel/transformers/tvd.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from liger_kernel.ops.tvd import LigerTVDLossFunction
+from liger_kernel.ops import LigerTVDLossFunction
 
 
 class LigerTVDLoss(nn.Module):