Implement RMSNORM in cuTile and cuTeDSL

src/liger_kernel/ops/cutedsl/__init__.py

@@ -0,0 +1,17 @@
+"""
+cuTeDSL backend for Liger-Kernel.
+
+This backend is opt-in only and currently ships RMSNorm parity first.
+Enable it with ``LIGER_KERNEL_IMPL=cutedsl``.
+"""
+
+from liger_kernel.ops.backends.registry import ImplInfo
+from liger_kernel.ops.backends.registry import register_impl
+
+register_impl(
+    ImplInfo(
+        name="cutedsl",
+        devices=("cuda",),
+        module_path=f"{__name__}.ops",  # liger_kernel.ops.cutedsl.ops
+    )
+)

src/liger_kernel/ops/cutedsl/ops/__init__.py

@@ -0,0 +1,22 @@
+"""
+cuTeDSL-specific operator implementations.
+"""
+
+try:
+    # Current parity implementation shares the same CUDA DSL runtime.
+    import cuda.tile as ct  # noqa: F401
+except ImportError as exc:
+    raise ImportError(
+        "cuTeDSL backend currently requires the CUDA tile runtime. Install with `pip install cuda-tile` "
+        "or `pip install 'cuda-tile[tileiras]'`."
+    ) from exc
+
+from liger_kernel.ops.cutedsl.ops.rms_norm import LigerRMSNormFunction
+from liger_kernel.ops.cutedsl.ops.rms_norm import rms_norm_backward
+from liger_kernel.ops.cutedsl.ops.rms_norm import rms_norm_forward
+
+__all__ = [
+    "LigerRMSNormFunction",
+    "rms_norm_backward",
+    "rms_norm_forward",
+]

src/liger_kernel/ops/cutedsl/ops/rms_norm.py

@@ -0,0 +1,16 @@
+"""
+RMSNorm implementation for the cuTeDSL backend.
+
+This module reuses the cuTile RMSNorm kernels to ensure behavior parity while
+the dedicated cuTeDSL kernel implementation is being developed.
+"""
+
+from liger_kernel.ops.cutile.ops.rms_norm import LigerRMSNormFunction
+from liger_kernel.ops.cutile.ops.rms_norm import rms_norm_backward
+from liger_kernel.ops.cutile.ops.rms_norm import rms_norm_forward
+
+__all__ = [
+    "LigerRMSNormFunction",
+    "rms_norm_forward",
+    "rms_norm_backward",
+]

src/liger_kernel/ops/cutile/ops/__init__.py

@@ -27,6 +27,9 @@
 from liger_kernel.ops.cutile.ops.layer_norm import LigerLayerNormFunction
 from liger_kernel.ops.cutile.ops.layer_norm import layer_norm_backward
 from liger_kernel.ops.cutile.ops.layer_norm import layer_norm_forward
+from liger_kernel.ops.cutile.ops.rms_norm import LigerRMSNormFunction
+from liger_kernel.ops.cutile.ops.rms_norm import rms_norm_backward
+from liger_kernel.ops.cutile.ops.rms_norm import rms_norm_forward
 
 __all__ = [
     "LigerCrossEntropyFunction",
@@ -44,4 +47,7 @@
     "LigerLayerNormFunction",
     "layer_norm_backward",
     "layer_norm_forward",
+    "LigerRMSNormFunction",
+    "rms_norm_backward",
+    "rms_norm_forward",
 ]