liger-kernel-nightly 0.5.6.dev20250407214804__py3-none-any.whl → 0.5.6.dev20250408194537__py3-none-any.whl

liger_kernel/ops/layer_norm.py

@@ -154,6 +154,11 @@ def layer_norm_forward(X, W, B, eps):
             f"must match weight size (W.shape[0]={W.shape[0]})"
         )
 
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args["grf_mode"] = "large"
+
     _layer_norm_forward_kernel[(n_rows,)](
         Y,
         Y.stride(0),
@@ -171,6 +176,7 @@ def layer_norm_forward(X, W, B, eps):
         eps,
         BLOCK_SIZE=BLOCK_SIZE,
         num_warps=num_warps,
+        **kernel_args,  # XPU-specific optimization
     )
     return Y.view(*shape), X, Mean, RSTD, BLOCK_SIZE, num_warps
 
@@ -185,7 +191,7 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
     if X.device.type == "cuda":
         sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
     elif X.device.type == "xpu":
-        sm_count = torch.xpu.get_device_properties(X.device).gpu_subslice_count
+        sm_count = torch.xpu.get_device_properties(X.device).gpu_eu_count
 
     DX = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
     _DW = torch.empty((sm_count, n_cols), dtype=W.dtype, device=W.device)
@@ -208,6 +214,12 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
         if X.dtype == torch.float16
         else tl.float32  # fallback to float32 for other types
     )
+
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args.update({"grf_mode": "large", "num_warps": 32, "num_stages": 4})
+
     _layer_norm_backward_kernel[grid](
         X,
         W,
@@ -227,6 +239,7 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
         rows_per_program,
         BLOCK_SIZE=BLOCK_SIZE,
         dtype=triton_dtype,
+        **kernel_args,  # XPU-specific optimization
     )
 
     DW = _DW.sum(dim=0).to(W.dtype)

liger_kernel/ops/rms_norm.py

@@ -223,6 +223,10 @@ def rms_norm_forward(X, W, eps, offset, casting_mode):
     # Check constraints.
     assert X.shape[1] == W.shape[0], "Incompatible hidden size dimension between tensor1.shape[1] and tensor2.shape[0]"
 
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args["grf_mode"] = "large"
     _rms_norm_forward_kernel[(n_rows,)](
         Y,
         Y.stride(0),
@@ -238,6 +242,7 @@ def rms_norm_forward(X, W, eps, offset, casting_mode):
         casting_mode,
         BLOCK_SIZE=BLOCK_SIZE,
         num_warps=num_warps,
+        **kernel_args,  # XPU-specific optimization
     )
     return Y.view(*shape), X, RSTD, BLOCK_SIZE, num_warps, casting_mode
 
@@ -252,7 +257,7 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
     if X.device.type == "cuda":
         sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
     elif X.device.type == "xpu":
-        sm_count = torch.xpu.get_device_properties(X.device).gpu_subslice_count
+        sm_count = torch.xpu.get_device_properties(X.device).gpu_eu_count
 
     # fp32 for numerical stability especially.
     _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
@@ -267,6 +272,11 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
     else:
         dX = torch.zeros_like(dY)
 
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args["grf_mode"] = "large"
+
     _rms_norm_backward_kernel[grid](
         dY,
         dY.stride(0),
@@ -288,6 +298,7 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
         casting_mode,
         BLOCK_SIZE=BLOCK_SIZE,
         num_warps=num_warps,
+        **kernel_args,  # XPU-specific optimization
     )
     dX = dX.view(*shape)
     dW = _dW.sum(dim=0).to(W.dtype)

liger_kernel/transformers/__init__.py

@@ -1,4 +1,6 @@
-from liger_kernel.transformers.auto_model import AutoLigerKernelForCausalLM  # noqa: F401
+import importlib
+
+# Always-safe imports (independent of 'transformers')
 from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss  # noqa: F401
 from liger_kernel.transformers.dyt import LigerDyT  # noqa: F401
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss  # noqa: F401
@@ -6,27 +8,113 @@ from liger_kernel.transformers.fused_linear_jsd import LigerFusedLinearJSD  # no
 from liger_kernel.transformers.geglu import LigerGEGLUMLP  # noqa: F401
 from liger_kernel.transformers.jsd import LigerJSD  # noqa: F401
 from liger_kernel.transformers.layer_norm import LigerLayerNorm  # noqa: F401
-from liger_kernel.transformers.monkey_patch import _apply_liger_kernel  # noqa: F401
-from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma2  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3_text  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_granite  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llava  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mistral  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mixtral  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mllama  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_olmo2  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_paligemma  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_phi3  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_5_vl  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_vl  # noqa: F401
 from liger_kernel.transformers.rms_norm import LigerRMSNorm  # noqa: F401
 from liger_kernel.transformers.rope import liger_rotary_pos_emb  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerSwiGLUMLP  # noqa: F401
 from liger_kernel.transformers.tvd import LigerTVDLoss  # noqa: F401
+
+# Check if 'transformers' is installed
+try:
+    import transformers  # noqa: F401
+
+    _TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    _TRANSFORMERS_AVAILABLE = False
+
+
+def is_transformers_available() -> bool:
+    """
+    Returns True if the 'transformers' package is available.
+    Useful for conditional logic in downstream code.
+    """
+    return _TRANSFORMERS_AVAILABLE
+
+
+def __getattr__(name: str):
+    """
+    Handles lazy access to transformer-dependent attributes.
+    If 'transformers' is not installed, raises a user-friendly ImportError.
+    """
+    if not _TRANSFORMERS_AVAILABLE:
+        raise ImportError(
+            f"The attribute '{name}' requires the 'transformers' library, which is not installed.\n"
+            f"Please install it with `pip install transformers` to use this functionality."
+        )
+
+    if name == "AutoLigerKernelForCausalLM":
+        module = importlib.import_module("liger_kernel.transformers.auto_model")
+        return getattr(module, name)
+
+    monkey_patch_symbols = {
+        "_apply_liger_kernel",
+        "_apply_liger_kernel_to_instance",
+        "apply_liger_kernel_to_gemma",
+        "apply_liger_kernel_to_gemma2",
+        "apply_liger_kernel_to_gemma3",
+        "apply_liger_kernel_to_gemma3_text",
+        "apply_liger_kernel_to_granite",
+        "apply_liger_kernel_to_llama",
+        "apply_liger_kernel_to_llava",
+        "apply_liger_kernel_to_mistral",
+        "apply_liger_kernel_to_mixtral",
+        "apply_liger_kernel_to_mllama",
+        "apply_liger_kernel_to_olmo2",
+        "apply_liger_kernel_to_paligemma",
+        "apply_liger_kernel_to_phi3",
+        "apply_liger_kernel_to_qwen2",
+        "apply_liger_kernel_to_qwen2_5_vl",
+        "apply_liger_kernel_to_qwen2_vl",
+    }
+
+    if name in monkey_patch_symbols:
+        module = importlib.import_module("liger_kernel.transformers.monkey_patch")
+        return getattr(module, name)
+
+    raise AttributeError(f"module {__name__} has no attribute {name}")
+
+
+# Shared symbols in all environments
+__all__ = [
+    "is_transformers_available",
+    "LigerCrossEntropyLoss",
+    "LigerDyT",
+    "LigerFusedLinearCrossEntropyLoss",
+    "LigerFusedLinearJSD",
+    "LigerGEGLUMLP",
+    "LigerJSD",
+    "LigerLayerNorm",
+    "LigerRMSNorm",
+    "liger_rotary_pos_emb",
+    "LigerBlockSparseTop2MLP",
+    "LigerPhi3SwiGLUMLP",
+    "LigerSwiGLUMLP",
+    "LigerTVDLoss",
+]
+
+# Add transformer-dependent symbols only if available
+if _TRANSFORMERS_AVAILABLE:
+    __all__.extend(
+        [
+            "AutoLigerKernelForCausalLM",
+            "_apply_liger_kernel",
+            "_apply_liger_kernel_to_instance",
+            "apply_liger_kernel_to_gemma",
+            "apply_liger_kernel_to_gemma2",
+            "apply_liger_kernel_to_gemma3",
+            "apply_liger_kernel_to_gemma3_text",
+            "apply_liger_kernel_to_granite",
+            "apply_liger_kernel_to_llama",
+            "apply_liger_kernel_to_llava",
+            "apply_liger_kernel_to_mistral",
+            "apply_liger_kernel_to_mixtral",
+            "apply_liger_kernel_to_mllama",
+            "apply_liger_kernel_to_olmo2",
+            "apply_liger_kernel_to_paligemma",
+            "apply_liger_kernel_to_phi3",
+            "apply_liger_kernel_to_qwen2",
+            "apply_liger_kernel_to_qwen2_5_vl",
+            "apply_liger_kernel_to_qwen2_vl",
+        ]
+    )

{liger_kernel_nightly-0.5.6.dev20250407214804.dist-info → liger_kernel_nightly-0.5.6.dev20250408194537.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.6.dev20250407214804
+Version: 0.5.6.dev20250408194537
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

{liger_kernel_nightly-0.5.6.dev20250407214804.dist-info → liger_kernel_nightly-0.5.6.dev20250408194537.dist-info}/RECORD

@@ -24,16 +24,16 @@ liger_kernel/ops/geglu.py,sha256=axGvCIvlBzuluoAIrWTsp2iZM4BFKNInkPov8YVvH9E,412
 liger_kernel/ops/group_norm.py,sha256=qD4D4lSjSgVtO52EBNLC2iTseALRgPgqXE50U2woggk,10837
 liger_kernel/ops/jsd.py,sha256=onHp5T3MbvJaVz5Vup7Ww6EQp_HTaZeayTjJk6FgQMY,7042
 liger_kernel/ops/kl_div.py,sha256=NkG7D6_DnPBzr-ohhYiQbRBnq_fbGmpn5UU7y0UBKQo,8420
-liger_kernel/ops/layer_norm.py,sha256=6roQjioyg-9O2qLPV8nL4U0-5UH80tdzOMTWwjvDnn8,7961
+liger_kernel/ops/layer_norm.py,sha256=vWCyOm-F2GMAilB-ozJcFeUQQLCJoTE_uiXq-_0uYuI,8356
 liger_kernel/ops/qwen2vl_mrope.py,sha256=3GExhYpLgB4VUtyZyjRk8XjEur3W4EWF6HQ67ML5vBU,8481
-liger_kernel/ops/rms_norm.py,sha256=PWLJcdIKU5e-8BuYFHd9Cqlq6wmr6fUXKi9zQD4LetU,11727
+liger_kernel/ops/rms_norm.py,sha256=PP27OIBmV9By63i13jot9ylDowW0nuxY_JFIkaPLgL4,12078
 liger_kernel/ops/rope.py,sha256=ofmBOkUpZZO-Q8Z5B_LOFYYLD-YT-8WnJ4vGOrDYouI,8943
 liger_kernel/ops/swiglu.py,sha256=KmgMjaJQnbLLgZn2nEpbwHU_xpnYRweCyrLQSVvM1vA,3015
 liger_kernel/ops/tvd.py,sha256=FHJtLQI95ijqgg9UtaHpMAjSCiPxB6CduPwPMcGxelc,6405
 liger_kernel/ops/utils.py,sha256=uoFKQqo-34N2TWQNvXMFywqGiOMMXNEVBxVojzlUAa0,3836
 liger_kernel/ops/experimental/embedding.py,sha256=tolj3tItkzpSb30zWqDN2_yX4ectflaQ8HMyKyFIQc8,4172
 liger_kernel/ops/experimental/mm_int8int2.py,sha256=TrS9lpwekrik_w5qE7AhMJD1bcq-OidjtbsW80oZ6IM,13314
-liger_kernel/transformers/__init__.py,sha256=23EWILU5JpoZjEx_yZCQMe6lprEsBJGHVt5GQBbwptg,2811
+liger_kernel/transformers/__init__.py,sha256=d02IvvSSxSzLenLtexzi7pHUCfl0BGGV0TxOErzb1VE,4454
 liger_kernel/transformers/auto_model.py,sha256=0qCTRZt280Bj_LcFdzo9hlaR-BWNazawXOGgoCZjgEg,1545
 liger_kernel/transformers/cross_entropy.py,sha256=z3KTWQnFxr_IZaVjtYt0ZNEWQdDdYThN35xWkHlDGH0,1683
 liger_kernel/transformers/dyt.py,sha256=QMqqc14pkE0WhpRZvapfnNAun-6C0C_tHExL2ZJuCUA,648
@@ -74,9 +74,9 @@ liger_kernel/transformers/trainer/__init__.py,sha256=p7yQfklV8-467qSz_ZMimkbDF7H
 liger_kernel/transformers/trainer/orpo_trainer.py,sha256=pdekW7l6Qg_aqa5SYKYlSWUF8m3lkOFvFLcIMEHrz9s,8338
 liger_kernel/triton/__init__.py,sha256=qCiCamzCRv6lpV8IqpAc9YMdNKC7GKurClWceQPnlis,92
 liger_kernel/triton/monkey_patch.py,sha256=Rd0hUHAzDkFfHvnX7-PBaNK5EKnZhtfM_h-fgQH9HPY,1568
-liger_kernel_nightly-0.5.6.dev20250407214804.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.5.6.dev20250407214804.dist-info/METADATA,sha256=0lQVqhPNaqGVZvOrb6MxIp2eP7IYoABa4llfB8Ua868,23297
-liger_kernel_nightly-0.5.6.dev20250407214804.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.5.6.dev20250407214804.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-liger_kernel_nightly-0.5.6.dev20250407214804.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.5.6.dev20250407214804.dist-info/RECORD,,
+liger_kernel_nightly-0.5.6.dev20250408194537.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.5.6.dev20250408194537.dist-info/METADATA,sha256=QqNEMLqg80zH5iTfcmAvYvM3ukixN6OSByeATy0BImM,23297
+liger_kernel_nightly-0.5.6.dev20250408194537.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.5.6.dev20250408194537.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+liger_kernel_nightly-0.5.6.dev20250408194537.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.5.6.dev20250408194537.dist-info/RECORD,,