liger-kernel-nightly 0.6.3.dev20251121202601__tar.gz → 0.6.4.dev20251229100549__tar.gz

{liger_kernel_nightly-0.6.3.dev20251121202601 → liger_kernel_nightly-0.6.4.dev20251229100549}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.6.3.dev20251121202601
+Version: 0.6.4.dev20251229100549
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -79,8 +79,8 @@ Requires-Dist: torchvision>=0.20; extra == "dev"
             </a>
         </td>
         <td style="padding: 10px;">
-            <a href="https://discord.gg/gpumode">
-                <img src="https://dcbadge.limes.pink/api/server/gpumode?style=flat" alt="Join Our Discord">
+            <a href="https://discord.gg/X4MaxPgA">
+                <img src="https://dcbadge.limes.pink/api/server/https://discord.gg/X4MaxPgA?style=flat" alt="Join Our Discord">
             </a>
         </td>
     </tr>
@@ -95,6 +95,7 @@ Requires-Dist: torchvision>=0.20; extra == "dev"
 <details>
   <summary>Latest News 🔥</summary>
 
+  - [2025/12/19] We announced a liger kernel discord channel at https://discord.gg/X4MaxPgA; We will be hosting Liger Kernel x Triton China Meetup in mid of January 2026
   - [2025/03/06] We release a joint blog post on TorchTune × Liger - [Peak Performance, Minimized Memory: Optimizing torchtune’s performance with torch.compile & Liger Kernel](https://pytorch.org/blog/peak-performance-minimized-memory/)
   - [2024/12/11] We release [v0.5.0](https://github.com/linkedin/Liger-Kernel/releases/tag/v0.5.0): 80% more memory efficient post training losses (DPO, ORPO, CPO, etc)!
   - [2024/12/5] We release LinkedIn Engineering Blog - [Liger-Kernel: Empowering an open source ecosystem of Triton Kernels for Efficient LLM Training](https://www.linkedin.com/blog/engineering/open-source/liger-kernel-open-source-ecosystem-for-efficient-llm-training)
@@ -113,6 +114,8 @@ We've also added optimized Post-Training kernels that deliver **up to 80% memory
 
 You can view the documentation site for additional installation, usage examples, and API references:https://linkedin.github.io/Liger-Kernel/
 
+You can view the Liger Kernel Technical Report: https://openreview.net/forum?id=36SjAIT42G
+
 ## Supercharge Your Model with Liger Kernel
 
 ![Banner](https://raw.githubusercontent.com/linkedin/Liger-Kernel/main/docs/images/banner.GIF)
@@ -312,6 +315,7 @@ loss.backward()
 | OLMo2   | `liger_kernel.transformers.apply_liger_kernel_to_olmo2`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | Olmo3   | `liger_kernel.transformers.apply_liger_kernel_to_olmo3`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | GLM-4   | `liger_kernel.transformers.apply_liger_kernel_to_glm4`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
+| GPT-OSS   | `liger_kernel.transformers.apply_liger_kernel_to_gpt_oss`     | RoPE, RMSNorm, CrossEntropyLoss, FusedLinearCrossEntropy |
 | InternVL3   | `liger_kernel.transformers.apply_liger_kernel_to_internvl`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | HunyuanV1   | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_dense`    |  RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy       |
 | HunyuanV1 MoE | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_moe` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy       |
@@ -441,3 +445,5 @@ url={https://openreview.net/forum?id=36SjAIT42G}
         ↑ Back to Top ↑
     </a>
 </p>
+
+

{liger_kernel_nightly-0.6.3.dev20251121202601 → liger_kernel_nightly-0.6.4.dev20251229100549}/README.md

@@ -31,8 +31,8 @@
             </a>
         </td>
         <td style="padding: 10px;">
-            <a href="https://discord.gg/gpumode">
-                <img src="https://dcbadge.limes.pink/api/server/gpumode?style=flat" alt="Join Our Discord">
+            <a href="https://discord.gg/X4MaxPgA">
+                <img src="https://dcbadge.limes.pink/api/server/https://discord.gg/X4MaxPgA?style=flat" alt="Join Our Discord">
             </a>
         </td>
     </tr>
@@ -47,6 +47,7 @@
 <details>
   <summary>Latest News 🔥</summary>
 
+  - [2025/12/19] We announced a liger kernel discord channel at https://discord.gg/X4MaxPgA; We will be hosting Liger Kernel x Triton China Meetup in mid of January 2026
   - [2025/03/06] We release a joint blog post on TorchTune × Liger - [Peak Performance, Minimized Memory: Optimizing torchtune’s performance with torch.compile & Liger Kernel](https://pytorch.org/blog/peak-performance-minimized-memory/)
   - [2024/12/11] We release [v0.5.0](https://github.com/linkedin/Liger-Kernel/releases/tag/v0.5.0): 80% more memory efficient post training losses (DPO, ORPO, CPO, etc)!
   - [2024/12/5] We release LinkedIn Engineering Blog - [Liger-Kernel: Empowering an open source ecosystem of Triton Kernels for Efficient LLM Training](https://www.linkedin.com/blog/engineering/open-source/liger-kernel-open-source-ecosystem-for-efficient-llm-training)
@@ -65,6 +66,8 @@ We've also added optimized Post-Training kernels that deliver **up to 80% memory
 
 You can view the documentation site for additional installation, usage examples, and API references:https://linkedin.github.io/Liger-Kernel/
 
+You can view the Liger Kernel Technical Report: https://openreview.net/forum?id=36SjAIT42G
+
 ## Supercharge Your Model with Liger Kernel
 
 ![Banner](https://raw.githubusercontent.com/linkedin/Liger-Kernel/main/docs/images/banner.GIF)
@@ -264,6 +267,7 @@ loss.backward()
 | OLMo2   | `liger_kernel.transformers.apply_liger_kernel_to_olmo2`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | Olmo3   | `liger_kernel.transformers.apply_liger_kernel_to_olmo3`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | GLM-4   | `liger_kernel.transformers.apply_liger_kernel_to_glm4`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
+| GPT-OSS   | `liger_kernel.transformers.apply_liger_kernel_to_gpt_oss`     | RoPE, RMSNorm, CrossEntropyLoss, FusedLinearCrossEntropy |
 | InternVL3   | `liger_kernel.transformers.apply_liger_kernel_to_internvl`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | HunyuanV1   | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_dense`    |  RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy       |
 | HunyuanV1 MoE | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_moe` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy       |
@@ -393,3 +397,5 @@ url={https://openreview.net/forum?id=36SjAIT42G}
         ↑ Back to Top ↑
     </a>
 </p>
+
+

{liger_kernel_nightly-0.6.3.dev20251121202601 → liger_kernel_nightly-0.6.4.dev20251229100549}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.6.3.dev20251121202601"
+version = "0.6.4.dev20251229100549"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

{liger_kernel_nightly-0.6.3.dev20251121202601 → liger_kernel_nightly-0.6.4.dev20251229100549}/setup.py

@@ -24,6 +24,8 @@ def get_default_dependencies():
         return [
             "torch>=2.6.0",
         ]
+    elif platform == "npu":
+        return ["torch_npu==2.7.1", "triton-ascend"]
 
 
 def get_optional_dependencies():
@@ -67,7 +69,21 @@ def is_xpu_available():
     return False
 
 
-def get_platform() -> Literal["cuda", "rocm", "cpu", "xpu"]:
+def is_ascend_available() -> bool:
+    """Best-effort Ascend detection.
+
+    Checks for common Ascend environment variables and a possible `npu-smi`
+    utility if present.
+    """
+    try:
+        subprocess.run(["npu-smi", "info"], check=True)
+        return True
+    except (subprocess.SubprocessError, FileNotFoundError):
+        pass
+    return False
+
+
+def get_platform() -> Literal["cuda", "rocm", "cpu", "xpu", "npu"]:
     """
     Detect whether the system has NVIDIA or AMD GPU without torch dependency.
     """
@@ -86,6 +102,9 @@ def get_platform() -> Literal["cuda", "rocm", "cpu", "xpu"]:
             if is_xpu_available():
                 print("Intel GPU detected")
                 return "xpu"
+            elif is_ascend_available():
+                print("Ascend NPU detected")
+                return "npu"
             else:
                 print("No GPU detected")
                 return "cpu"

{liger_kernel_nightly-0.6.3.dev20251121202601 → liger_kernel_nightly-0.6.4.dev20251229100549}/src/liger_kernel/chunked_loss/cosine_similarity_loss.py

@@ -9,7 +9,13 @@ from liger_kernel.chunked_loss.fused_linear_distillation import LigerFusedLinear
 
 class LigerFusedLinearCosineSimilarityFunction(LigerFusedLinearDistillationBase):
     @staticmethod
-    def distillation_loss_fn(student_logits, teacher_logits, beta=1.0):
+    def distillation_loss_fn(
+        student_logits,
+        teacher_logits,
+        target=None,
+        ignore_index=None,
+        beta=1.0,
+    ):
         """
         Compute Cosine loss (Cosine Similarity Loss).
         Args:

{liger_kernel_nightly-0.6.3.dev20251121202601 → liger_kernel_nightly-0.6.4.dev20251229100549}/src/liger_kernel/chunked_loss/fused_linear_distillation.py

@@ -13,6 +13,8 @@ class LigerFusedLinearDistillationBase(torch.autograd.Function):
     def distillation_loss_fn(
         student_logits,
         teacher_logits,
+        target=None,
+        ignore_index=None,
     ):
         """
         Compute distillation loss.
@@ -132,10 +134,15 @@ class LigerFusedLinearDistillationBase(torch.autograd.Function):
             )
             student_logits_chunk = torch.cat([student_logits_chunk, pad_tensor], dim=-1)
 
-        hard_loss /= full_target.shape[0]
+        num_valid_tokens = (full_target != ignore_index).sum()
+        num_valid_tokens = num_valid_tokens.clamp_min(1)  # to avoid division by zero
 
-        soft_loss = distillation_loss_fn(student_logits_chunk, teacher_logits_chunk, **loss_kwargs)
-        soft_loss /= full_target.shape[0]
+        hard_loss /= num_valid_tokens
+
+        soft_loss = distillation_loss_fn(
+            student_logits_chunk, teacher_logits_chunk, target=target_chunk, ignore_index=ignore_index, **loss_kwargs
+        )
+        soft_loss /= num_valid_tokens
 
         loss = weight_hard_loss * hard_loss + weight_soft_loss * soft_loss
         return loss, (soft_loss, hard_loss, student_logits_chunk, teacher_logits_chunk)

{liger_kernel_nightly-0.6.3.dev20251121202601 → liger_kernel_nightly-0.6.4.dev20251229100549}/src/liger_kernel/chunked_loss/jsd_loss.py

@@ -11,35 +11,50 @@ from liger_kernel.chunked_loss.fused_linear_distillation import LigerFusedLinear
 
 class LigerFusedLinearJSDFunction(LigerFusedLinearDistillationBase):
     @staticmethod
-    def distillation_loss_fn(student_logits, teacher_logits, beta=0.5):
+    def distillation_loss_fn(student_logits, teacher_logits, beta=0.5, target=None, ignore_index=-100):
         """
         Compute JSD loss (Jensen-Shannon Divergence Loss).
         Args:
             student_logits (torch.Tensor): Logits of student tokens. Shape: (batch_size * seq_len,).
             teacher_logits (torch.Tensor): Logits of teacher tokens. Shape: (batch_size * seq_len,).
             beta (float): Coefficient beta of generalized JSD in the interval [0, 1]. Default: `0.5`.
+            target (torch.Tensor): Target labels for masking. Shape: (chunk_size,).
+            ignore_index (int): Index to ignore in loss computation.
         Returns:
             torch.Tensor: Jensen-Shannon Divergence loss
+        Note:
+            - Uses reduction="none" to preserve per-token losses for masking
+            - KL divergence requires summing over vocab dimension (not mean)
+            - Masking excludes padding/prompt tokens from loss computation
         """
         student_log_probs = F.log_softmax(student_logits, dim=-1)
         teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)
 
         if beta == 0:
-            jsd_loss = F.kl_div(student_log_probs, teacher_log_probs, reduction="sum", log_target=True)
+            jsd_loss = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True)
         elif beta == 1:
-            jsd_loss = F.kl_div(teacher_log_probs, student_log_probs, reduction="sum", log_target=True)
+            jsd_loss = F.kl_div(teacher_log_probs, student_log_probs, reduction="none", log_target=True)
         else:
             # Compute probabilities (only required for mean calculation)
             log_mean_probs = torch.logsumexp(
                 torch.stack([student_log_probs + math.log(1 - beta), teacher_log_probs + math.log(beta)], dim=0), dim=0
             )
 
-            student_kl = F.kl_div(log_mean_probs, student_log_probs, reduction="sum", log_target=True)
-            teacher_kl = F.kl_div(log_mean_probs, teacher_log_probs, reduction="sum", log_target=True)
+            student_kl = F.kl_div(log_mean_probs, student_log_probs, reduction="none", log_target=True)
+            teacher_kl = F.kl_div(log_mean_probs, teacher_log_probs, reduction="none", log_target=True)
 
             # JSD is the weighted average of the KL divergences
             jsd_loss = beta * teacher_kl + (1 - beta) * student_kl
-        return jsd_loss
+
+        # Sum over vocab dimension (KL divergence definition)
+        jsd_loss = jsd_loss.sum(dim=-1)  # (chunk_size,)
+
+        # Apply ignore_index mask
+        if target is not None:
+            mask = target != ignore_index
+            jsd_loss = jsd_loss.masked_fill(~mask, 0.0)
+
+        return jsd_loss.sum()
 
     @classmethod
     def forward(

liger_kernel_nightly-0.6.4.dev20251229100549/src/liger_kernel/ops/__init__.py

@@ -0,0 +1,141 @@
+"""
+Liger-Kernel operators with automatic vendor-specific replacement.
+
+This module provides two ways to import operators:
+
+1. Import from this package (recommended for Function classes):
+       from liger_kernel.ops import LigerGELUMulFunction
+
+   This automatically uses vendor-specific implementation if available.
+
+2. Import from submodules (for kernel functions or specific access):
+       from liger_kernel.ops.geglu import geglu_forward, geglu_backward
+
+   This always uses the default implementation (no auto-replacement).
+
+The replacement mechanism:
+1. Default implementations are imported from individual modules (e.g., geglu.py)
+2. On module load, device is detected via infer_device()
+3. If running on a supported vendor device (npu, xpu, etc.), the default
+   implementations are replaced with vendor-specific ones
+4. All subsequent imports from this package get the replaced versions
+
+Note: Direct imports from submodules (e.g., from liger_kernel.ops.geglu import ...)
+      are NOT affected by the replacement mechanism.
+"""
+
+# =============================================================================
+# Import default implementations
+# Both Function classes and kernel functions are imported here.
+# All of these can be replaced by vendor-specific implementations.
+# =============================================================================
+
+from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction  # noqa: F401
+from liger_kernel.ops.cross_entropy import cross_entropy_backward  # noqa: F401
+from liger_kernel.ops.cross_entropy import cross_entropy_forward  # noqa: F401
+from liger_kernel.ops.dyt import LigerDyTFunction  # noqa: F401
+from liger_kernel.ops.experimental.embedding import LigerEmbeddingFunction  # noqa: F401
+from liger_kernel.ops.fused_add_rms_norm import LigerFusedAddRMSNormFunction  # noqa: F401
+from liger_kernel.ops.fused_add_rms_norm import fused_add_rms_norm_backward  # noqa: F401
+from liger_kernel.ops.fused_add_rms_norm import fused_add_rms_norm_forward  # noqa: F401
+from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction  # noqa: F401
+from liger_kernel.ops.fused_linear_cross_entropy import fused_linear_cross_entropy_backward  # noqa: F401
+from liger_kernel.ops.fused_linear_cross_entropy import fused_linear_cross_entropy_forward  # noqa: F401
+from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction  # noqa: F401
+from liger_kernel.ops.fused_linear_jsd import fused_linear_jsd_backward  # noqa: F401
+from liger_kernel.ops.fused_linear_jsd import fused_linear_jsd_forward  # noqa: F401
+from liger_kernel.ops.fused_neighborhood_attention import LigerFusedNeighborhoodAttentionFunction  # noqa: F401
+from liger_kernel.ops.geglu import LigerGELUMulFunction  # noqa: F401
+from liger_kernel.ops.geglu import geglu_backward  # noqa: F401
+from liger_kernel.ops.geglu import geglu_forward  # noqa: F401
+from liger_kernel.ops.group_norm import LigerGroupNormFunction  # noqa: F401
+from liger_kernel.ops.group_norm import group_norm_backward  # noqa: F401
+from liger_kernel.ops.group_norm import group_norm_forward  # noqa: F401
+from liger_kernel.ops.grpo_loss import GrpoLossFunction  # noqa: F401
+from liger_kernel.ops.jsd import LigerJSDFunction  # noqa: F401
+from liger_kernel.ops.jsd import jsd_backward  # noqa: F401
+from liger_kernel.ops.jsd import jsd_forward  # noqa: F401
+from liger_kernel.ops.kl_div import LigerKLDivLossFunction  # noqa: F401
+from liger_kernel.ops.layer_norm import LigerLayerNormFunction  # noqa: F401
+from liger_kernel.ops.layer_norm import layer_norm_backward  # noqa: F401
+from liger_kernel.ops.layer_norm import layer_norm_forward  # noqa: F401
+from liger_kernel.ops.llama4_rope import LigerLlama4RopeFunction  # noqa: F401
+from liger_kernel.ops.multi_token_attention import LigerMultiTokenAttentionFunction  # noqa: F401
+from liger_kernel.ops.poly_norm import LigerPolyNormFunction  # noqa: F401
+from liger_kernel.ops.poly_norm import poly_norm_backward  # noqa: F401
+from liger_kernel.ops.poly_norm import poly_norm_forward  # noqa: F401
+from liger_kernel.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction  # noqa: F401
+from liger_kernel.ops.rms_norm import LigerRMSNormFunction  # noqa: F401
+from liger_kernel.ops.rms_norm import rms_norm_backward  # noqa: F401
+from liger_kernel.ops.rms_norm import rms_norm_forward  # noqa: F401
+from liger_kernel.ops.rope import LigerRopeFunction  # noqa: F401
+from liger_kernel.ops.rope import rope_backward  # noqa: F401
+from liger_kernel.ops.rope import rope_forward  # noqa: F401
+from liger_kernel.ops.softmax import LigerSoftmaxFunction  # noqa: F401
+from liger_kernel.ops.sparsemax import LigerSparsemaxFunction  # noqa: F401
+from liger_kernel.ops.swiglu import LigerSiLUMulFunction  # noqa: F401
+from liger_kernel.ops.swiglu import swiglu_backward  # noqa: F401
+from liger_kernel.ops.swiglu import swiglu_forward  # noqa: F401
+from liger_kernel.ops.tiled_mlp import LigerTiledMLPFunction  # noqa: F401
+from liger_kernel.ops.tiled_mlp import apply_tiled_mlp  # noqa: F401
+from liger_kernel.ops.tvd import LigerTVDLossFunction  # noqa: F401
+
+# NOTE: __all__ is intentionally NOT defined.
+# - Import from this package (liger_kernel.ops) -> subject to vendor replacement
+# - Import from submodules (liger_kernel.ops.geglu) -> always use default implementation
+
+
+# =============================================================================
+# Vendor-specific replacement logic
+# =============================================================================
+
+
+def _replace_with_vendor_ops():
+    """
+    Replace/add vendor-specific operator implementations.
+
+    This function is called automatically on module load. It:
+    1. Detects the current device (cuda, npu, xpu, etc.)
+    2. Looks up the vendor for that device via VENDOR_REGISTRY
+    3. Loads and applies vendor-specific implementations
+
+    Vendor implementations should be placed in:
+        liger_kernel/ops/backends/_<vendor>/ops/
+
+    If the vendor module defines __all__, only those symbols are exported.
+    Otherwise, all public symbols (not starting with _) are auto-discovered.
+
+    Note: Vendor can both override existing ops AND add new vendor-specific ops.
+    """
+    from liger_kernel.ops.backends import get_vendor_for_device
+    from liger_kernel.utils import infer_device
+
+    device = infer_device()
+
+    # Look up vendor info for this device
+    vendor_info = get_vendor_for_device(device)
+    if vendor_info is None:
+        return
+
+    try:
+        import importlib
+
+        vendor_ops = importlib.import_module(vendor_info.module_path)
+
+        # Get names to export: use __all__ if defined, otherwise auto-discover
+        names_to_export = getattr(vendor_ops, "__all__", None)
+
+        if names_to_export is None:
+            # Auto-discover: find all public symbols (classes and functions)
+            names_to_export = [name for name in dir(vendor_ops) if not name.startswith("_")]
+
+        # Replace or add to this module's globals
+        for name in names_to_export:
+            globals()[name] = getattr(vendor_ops, name)
+
+    except ImportError:
+        # Vendor module not available, use default implementations
+        pass
+
+
+_replace_with_vendor_ops()

liger_kernel_nightly-0.6.4.dev20251229100549/src/liger_kernel/ops/backends/README.md

@@ -0,0 +1,151 @@
+# Adding a New Vendor Backend
+
+This directory contains vendor-specific operator implementations that automatically replace the default (CUDA) implementations when running on the corresponding device.
+
+## Concepts
+
+- **Vendor**: Chip manufacturer (e.g., `ascend`, `intel`, `nvidia`)
+- **Device**: Device type (e.g., `npu`, `xpu`, `cuda`)
+- **VendorInfo**: Defines the mapping between vendor and device
+
+## Directory Structure
+
+```
+backends/
+├── README.md          
+├── __init__.py         
+├── registry.py         # VendorInfo, register_vendor(), VENDOR_REGISTRY
+├── _ascend/            # Ascend (Huawei) vendor - supports NPU
+│   ├── __init__.py     # Registers VendorInfo for NPU
+│   └── ops/
+│       ├── __init__.py # Exports vendor-specific implementations
+│       └── geglu.py    # NPU-specific GEGLU implementation
+└── _<vendor>/          # Your new vendor backend
+    └── ...
+```
+
+## How It Works
+
+1. When `liger_kernel.ops.backends` is imported, it imports all vendor packages (e.g., `_ascend`)
+2. Each vendor's `__init__.py` calls `register_vendor()` to register itself
+3. When `liger_kernel.ops` is imported, `_replace_with_vendor_ops()` is called
+4. It detects the current device via `infer_device()` and looks up the vendor
+5. Vendor implementations replace/add to the `liger_kernel.ops` namespace
+
+## Adding a New Vendor
+
+### Step 1: Create Directory Structure
+
+```bash
+mkdir -p backends/_<vendor>/ops
+touch backends/_<vendor>/__init__.py
+touch backends/_<vendor>/ops/__init__.py
+```
+
+### Step 2: Register Your Vendor
+
+In `backends/_<vendor>/__init__.py`, register your vendor:
+
+```python
+"""
+<Vendor> backend for Liger-Kernel.
+"""
+
+from liger_kernel.ops.backends.registry import VendorInfo, register_vendor
+
+register_vendor(
+    VendorInfo(
+        vendor="<vendor>",
+        device="<device>",
+    )
+)
+```
+
+
+### Step 3: Ensure Device Detection Works
+
+Make sure `infer_device()` in `liger_kernel/utils.py` can detect your device:
+
+```python
+def infer_device():
+    if torch.cuda.is_available():
+        return "cuda"
+    if is_npu_available():
+        return "npu"
+    # Add your device detection here
+    if is_<device>_available():
+        return "<device>"
+    return "cpu"
+```
+
+### Step 4: Implement Vendor-Specific Operators
+
+Create operator files in `backends/_<vendor>/ops/`. For example, `geglu.py`:
+
+```python
+import torch
+
+class LigerGELUMulFunction(torch.autograd.Function):
+    """
+    Vendor-specific LigerGELUMulFunction implementation.
+    """
+    @staticmethod
+    def forward(ctx, a, b):
+        # Your vendor-specific forward implementation
+        ...
+
+    @staticmethod
+    def backward(ctx, dc):
+        # Your vendor-specific backward implementation
+        ...
+
+# Optional: vendor-specific kernel functions
+def geglu_forward_vendor(a, b):
+    ...
+
+def geglu_backward_vendor(a, b, dc):
+    ...
+```
+
+### Step 5: Export in `ops/__init__.py`
+
+In `backends/_<vendor>/ops/__init__.py`, export your implementations:
+
+```python
+"""
+<Vendor>-specific operator implementations.
+"""
+
+from .<module> import (
+    LigerGELUMulFunction,
+    geglu_forward_vendor as geglu_forward,   # Rename to match default API
+    geglu_backward_vendor as geglu_backward,
+)
+
+# Explicitly declare what to export (recommended)
+__all__ = [
+    "LigerGELUMulFunction",
+    "geglu_forward",
+    "geglu_backward",
+]
+```
+
+## Key Points
+
+### Incremental Override
+
+You **don't need to implement all operators**. Only implement the ones that require vendor-specific adaptations. Unimplemented operators will automatically fall back to the default (CUDA) implementation.
+
+### Vendor-Specific Additions
+
+Vendors can also **add new operators** that don't exist in the default implementation. These will be exported to `liger_kernel.ops` namespace for users to import.
+
+### Naming Convention
+
+- Use the **same class/function names** as the default implementations for overrides
+- This allows seamless replacement without changing user code
+- Use `as` imports to rename if your internal naming differs
+
+## Example: Ascend NPU Backend
+
+See `_ascend/` directory for a complete example of the Ascend NPU backend implementation.

liger_kernel_nightly-0.6.4.dev20251229100549/src/liger_kernel/ops/backends/__init__.py

@@ -0,0 +1,13 @@
+import importlib
+import pkgutil
+
+from liger_kernel.ops.backends.registry import VENDOR_REGISTRY  # noqa: F401
+from liger_kernel.ops.backends.registry import VendorInfo  # noqa: F401
+from liger_kernel.ops.backends.registry import get_vendor_for_device  # noqa: F401
+from liger_kernel.ops.backends.registry import register_vendor  # noqa: F401
+
+# Auto-import all _<vendor> subpackages to trigger registration
+# Each vendor's __init__.py calls register_vendor() when imported
+for _, modname, ispkg in pkgutil.iter_modules(__path__):
+    if ispkg and modname.startswith("_"):
+        importlib.import_module(f"{__name__}.{modname}")

liger_kernel_nightly-0.6.4.dev20251229100549/src/liger_kernel/ops/backends/_ascend/__init__.py

@@ -0,0 +1,5 @@
+from liger_kernel.ops.backends.registry import VendorInfo
+from liger_kernel.ops.backends.registry import register_vendor
+
+# Register Ascend vendor for NPU device
+register_vendor(VendorInfo(vendor="ascend", device="npu"))