issue/497 - rms norm interface

Author: wooway777

include/infinicore/ops.hpp

@@ -4,3 +4,4 @@
 #include "ops/matmul.hpp"
 #include "ops/ones.hpp"
 #include "ops/rearrange.hpp"
+#include "ops/rms_norm.hpp"

include/infinicore/ops/rms_norm.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class RMSNorm {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, float);
+    static void execute(Tensor y, Tensor x, Tensor weight, float epsilon = 1e-5f);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor rms_norm(Tensor x, Tensor weight, float epsilon = 1e-5f);
+void rms_norm_(Tensor y, Tensor x, Tensor weight, float epsilon = 1e-5f);
+} // namespace infinicore::op

python/infinicore/__init__.py

@@ -27,6 +27,7 @@
 from infinicore.ops.add import add
 from infinicore.ops.matmul import matmul
 from infinicore.ops.rearrange import rearrange
+from infinicore.ops.rms_norm import rms_norm
 from infinicore.tensor import (
     empty,
     from_blob,
@@ -67,6 +68,7 @@
     "add",
     "matmul",
     "rearrange",
+    "rms_norm",
     "empty",
     "from_blob",
     "ones",

python/infinicore/ops/rms_norm.py

@@ -0,0 +1,13 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def rms_norm(input, weight, epsilon=1e-5, *, out=None):
+    if out is None:
+        return Tensor(
+            _infinicore.rms_norm(input._underlying, weight._underlying, epsilon)
+        )
+
+    _infinicore.rms_norm_(
+        out._underlying, input._underlying, weight._underlying, epsilon
+    )

src/infinicore/op/rms_norm/rms_norm.cc

@@ -0,0 +1,24 @@
+#include "infinicore/ops/rms_norm.hpp"
+
+namespace infinicore::op {
+
+common::OpDispatcher<RMSNorm::schema> &RMSNorm::dispatcher() {
+    static common::OpDispatcher<RMSNorm::schema> dispatcher_;
+    return dispatcher_;
+};
+
+void RMSNorm::execute(Tensor y, Tensor x, Tensor weight, float epsilon) {
+    dispatcher().lookup(context::getDevice().getType())(y, x, weight, epsilon);
+}
+
+Tensor rms_norm(Tensor x, Tensor weight, float epsilon) {
+    auto y = Tensor::empty(x->shape(), x->dtype(), x->device());
+    rms_norm_(y, x, weight, epsilon);
+    return y;
+}
+
+void rms_norm_(Tensor y, Tensor x, Tensor weight, float epsilon) {
+    RMSNorm::execute(y, x, weight, epsilon);
+}
+
+} // namespace infinicore::op

src/infinicore/op/rms_norm/rms_norm_infiniop.cc

@@ -0,0 +1,52 @@
+#include "../../utils.hpp"
+#include "infinicore/common/hash.hpp"
+#include "infinicore/ops/common/cache.hpp"
+#include "infinicore/ops/rms_norm.hpp"
+#include <infiniop.h>
+
+namespace infinicore::op::rms_norm_impl::infiniop {
+
+thread_local common::OpCache<size_t, infiniopRMSNormDescriptor_t> caches(
+    100, // capacity
+    [](infiniopRMSNormDescriptor_t &desc) {
+        if (desc != nullptr) {
+            INFINICORE_CHECK_ERROR(infiniopDestroyRMSNormDescriptor(desc));
+            desc = nullptr;
+        }
+    });
+
+void calculate(Tensor y, Tensor x, Tensor weight, float epsilon) {
+    size_t seed = hash_combine(y, x, weight, epsilon);
+
+    auto device_type = context::getDevice().getType();
+    auto device_index = context::getDevice().getIndex();
+
+    auto &cache = caches.getCache(device_type, device_index);
+
+    auto desc_opt = cache.get(seed);
+    infiniopRMSNormDescriptor_t desc = nullptr;
+
+    if (!desc_opt) {
+        INFINICORE_CHECK_ERROR(infiniopCreateRMSNormDescriptor(
+            context::getInfiniopHandle(), &desc,
+            y->desc(), x->desc(), weight->desc(), epsilon));
+        cache.put(seed, desc);
+    } else {
+        desc = *desc_opt;
+    }
+
+    size_t workspace_size = 0;
+    INFINICORE_CHECK_ERROR(infiniopGetRMSNormWorkspaceSize(desc, &workspace_size));
+    std::shared_ptr<Memory> workspace = context::allocateMemory(workspace_size);
+
+    INFINICORE_CHECK_ERROR(infiniopRMSNorm(
+        desc, workspace->data(), workspace_size,
+        y->data(), x->data(), weight->data(), context::getStream()));
+}
+
+static bool registered = []() {
+    RMSNorm::dispatcher().registerAll(&calculate, false);
+    return true;
+}();
+
+} // namespace infinicore::op::rms_norm_impl::infiniop

src/infinicore/pybind11/ops.hpp

@@ -5,6 +5,7 @@
 #include "ops/add.hpp"
 #include "ops/matmul.hpp"
 #include "ops/rearrange.hpp"
+#include "ops/rms_norm.hpp"
 
 namespace py = pybind11;
 
@@ -14,6 +15,7 @@ inline void bind(py::module &m) {
     bind_add(m);
     bind_matmul(m);
     bind_rearrange(m);
+    bind_rms_norm(m);
 }
 
 } // namespace infinicore::ops

src/infinicore/pybind11/ops/rms_norm.hpp

@@ -0,0 +1,44 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+#include "infinicore/ops/rms_norm.hpp"
+
+namespace py = pybind11;
+
+namespace infinicore::ops {
+
+inline void bind_rms_norm(py::module &m) {
+    m.def("rms_norm",
+          &op::rms_norm,
+          py::arg("x"),
+          py::arg("weight"),
+          py::arg("epsilon") = 1e-5f,
+          R"doc(Root Mean Square Normalization.
+
+Args:
+    x: Input tensor
+    weight: Scale weights
+    epsilon: Small constant for numerical stability, default is 1e-5
+
+Returns:
+    Normalized tensor with same shape as input
+)doc");
+
+    m.def("rms_norm_",
+          &op::rms_norm_,
+          py::arg("y"),
+          py::arg("x"),
+          py::arg("weight"),
+          py::arg("epsilon") = 1e-5f,
+          R"doc(In-place Root Mean Square Normalization.
+
+Args:
+    y: Output tensor
+    x: Input tensor
+    weight: Scale weights
+    epsilon: Small constant for numerical stability, default is 1e-5
+)doc");
+}
+
+} // namespace infinicore::ops