feat: add fused rms_norm_rope operator (per-head RMSNorm + RoPE; CPU, NVIDIA)

include/infinicore/ops.hpp

@@ -73,6 +73,7 @@
 #include "ops/recurrent_gated_delta_rule.hpp"
 #include "ops/relu.hpp"
 #include "ops/rms_norm.hpp"
+#include "ops/rms_norm_rope.hpp"
 #include "ops/rope.hpp"
 #include "ops/rot.hpp"
 #include "ops/rotary_embedding.hpp"

include/infinicore/ops/rms_norm_rope.hpp

@@ -0,0 +1,23 @@
+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "../nn/rope.hpp"
+#include "../tensor.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(RMSNormRoPE, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, float, infinicore::nn::RoPE::Algo);
+
+// Internal: fused per-head RMSNorm + RoPE (full rotary), in-place on x
+// x: [num_tokens, num_heads, head_dim]
+void rms_norm_rope_(Tensor x,
+                    const Tensor &weight,
+                    const Tensor &pos_ids,
+                    const Tensor &sin_table,
+                    const Tensor &cos_table,
+                    float epsilon,
+                    infinicore::nn::RoPE::Algo algo);
+
+} // namespace infinicore::op

include/infiniop.h

@@ -128,6 +128,7 @@
 #include "infiniop/ops/recurrent_gated_delta_rule.h"
 #include "infiniop/ops/relu.h"
 #include "infiniop/ops/rms_norm.h"
+#include "infiniop/ops/rms_norm_rope.h"
 #include "infiniop/ops/rope.h"
 #include "infiniop/ops/rot.h"
 #include "infiniop/ops/rotg.h"

include/infiniop/ops/rms_norm_rope.h

@@ -0,0 +1,35 @@
+#ifndef __INFINIOP_RMS_NORM_ROPE_API_H__
+#define __INFINIOP_RMS_NORM_ROPE_API_H__
+
+#include "../operator_descriptor.h"
+#include "rope.h"
+
+typedef struct InfiniopDescriptor *infiniopRMSNormRoPEDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateRMSNormRoPEDescriptor(
+    infiniopHandle_t handle,
+    infiniopRMSNormRoPEDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t x_desc,
+    infiniopTensorDescriptor_t weight_desc,
+    infiniopTensorDescriptor_t pos_ids_desc,
+    infiniopTensorDescriptor_t sin_table_desc,
+    infiniopTensorDescriptor_t cos_table_desc,
+    float epsilon,
+    infiniopRoPEAlgo_t algo);
+
+__INFINI_C __export infiniStatus_t infiniopGetRMSNormRoPEWorkspaceSize(infiniopRMSNormRoPEDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopRMSNormRoPE(
+    infiniopRMSNormRoPEDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *x,
+    void const *weight,
+    void const *pos_ids,
+    void const *sin_table,
+    void const *cos_table,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyRMSNormRoPEDescriptor(infiniopRMSNormRoPEDescriptor_t desc);
+
+#endif

src/infinicore/ops/rms_norm_rope/rms_norm_rope.cc

@@ -0,0 +1,39 @@
+#include "infinicore/ops/rms_norm_rope.hpp"
+#include "../../utils.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(RMSNormRoPE);
+
+RMSNormRoPE::RMSNormRoPE(Tensor x,
+                         const Tensor &weight,
+                         const Tensor &pos_ids,
+                         const Tensor &sin_table,
+                         const Tensor &cos_table,
+                         float epsilon,
+                         infinicore::nn::RoPE::Algo algo) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, weight, pos_ids, sin_table, cos_table);
+    INFINICORE_GRAPH_OP_DISPATCH(x->device().getType(), x, weight, pos_ids, sin_table, cos_table, epsilon, algo);
+}
+
+void RMSNormRoPE::execute(Tensor x,
+                          const Tensor &weight,
+                          const Tensor &pos_ids,
+                          const Tensor &sin_table,
+                          const Tensor &cos_table,
+                          float epsilon,
+                          infinicore::nn::RoPE::Algo algo) {
+    INFINICORE_GRAPH_OP_RECORD_OR_RUN(RMSNormRoPE, x, weight, pos_ids, sin_table, cos_table, epsilon, algo);
+}
+
+void rms_norm_rope_(Tensor x,
+                    const Tensor &weight,
+                    const Tensor &pos_ids,
+                    const Tensor &sin_table,
+                    const Tensor &cos_table,
+                    float epsilon,
+                    infinicore::nn::RoPE::Algo algo) {
+    RMSNormRoPE::execute(x, weight, pos_ids, sin_table, cos_table, epsilon, algo);
+}
+
+} // namespace infinicore::op

src/infinicore/ops/rms_norm_rope/rms_norm_rope_infiniop.cc

@@ -0,0 +1,85 @@
+#include "infinicore/ops/rms_norm_rope.hpp"
+
+#include "../infiniop_impl.hpp"
+
+#include <infiniop/ops/rms_norm_rope.h>
+
+namespace infinicore::op::rms_norm_rope_impl::infiniop {
+
+INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, RMSNormRoPE, 100);
+
+struct PlannedMeta {
+    std::shared_ptr<Descriptor> descriptor;
+    graph::GraphTensor workspace;
+    graph::GraphTensor x;
+    graph::GraphTensor weight;
+    graph::GraphTensor pos;
+    graph::GraphTensor sin;
+    graph::GraphTensor cos;
+};
+
+static infiniopRoPEAlgo_t to_infiniop_algo(infinicore::nn::RoPE::Algo algo) {
+    switch (algo) {
+    case infinicore::nn::RoPE::Algo::GPT_J:
+        return INFINIOP_ROPE_ALGO_GPT_J;
+    case infinicore::nn::RoPE::Algo::GPT_NEOX:
+        return INFINIOP_ROPE_ALGO_GPT_NEOX;
+    default:
+        throw std::runtime_error("Unsupported RoPE algorithm");
+    }
+}
+
+void *plan(Tensor x,
+           const Tensor &weight,
+           const Tensor &pos,
+           const Tensor &sin,
+           const Tensor &cos,
+           float epsilon,
+           infinicore::nn::RoPE::Algo algo) {
+    auto infiniop_algo = to_infiniop_algo(algo);
+    size_t key = hash_combine(x, weight, pos, sin, cos, epsilon, static_cast<int>(infiniop_algo));
+
+    INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(
+        Descriptor, descriptor, RMSNormRoPE, key, x->desc(),
+        weight->desc(),
+        pos->desc(),
+        sin->desc(),
+        cos->desc(),
+        epsilon,
+        infiniop_algo);
+
+    INFINIOP_WORKSPACE_TENSOR(workspace, RMSNormRoPE, descriptor);
+    return new PlannedMeta{
+        descriptor,
+        graph::GraphTensor(workspace),
+        graph::GraphTensor(x),
+        graph::GraphTensor(weight),
+        graph::GraphTensor(pos),
+        graph::GraphTensor(sin),
+        graph::GraphTensor(cos)};
+}
+
+void run(void *planned_meta) {
+    auto *p = reinterpret_cast<PlannedMeta *>(planned_meta);
+
+    INFINICORE_CHECK_ERROR(
+        infiniopRMSNormRoPE(
+            p->descriptor->desc,
+            p->workspace->data(),
+            p->workspace->numel(),
+            p->x->data(),
+            p->weight->data(),
+            p->pos->data(),
+            p->sin->data(),
+            p->cos->data(),
+            context::getStream()));
+}
+
+void cleanup(void **planned_meta_ptr) {
+    delete *reinterpret_cast<PlannedMeta **>(planned_meta_ptr);
+    *planned_meta_ptr = nullptr;
+}
+
+INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(RMSNormRoPE, &plan, &run, &cleanup);
+
+} // namespace infinicore::op::rms_norm_rope_impl::infiniop

src/infiniop/ops/rms_norm_rope/cpu/rms_norm_rope_cpu.cc

@@ -0,0 +1,135 @@
+#include "rms_norm_rope_cpu.h"
+#include "../../../devices/cpu/common_cpu.h"
+
+namespace op::rms_norm_rope::cpu {
+
+Descriptor::~Descriptor() {}
+
+infiniStatus_t Descriptor::create(
+    infiniopHandle_t handle,
+    Descriptor **desc_ptr,
+    infiniopTensorDescriptor_t x_desc,
+    infiniopTensorDescriptor_t w_desc,
+    infiniopTensorDescriptor_t pos_desc,
+    infiniopTensorDescriptor_t sin_desc,
+    infiniopTensorDescriptor_t cos_desc,
+    float epsilon,
+    infiniopRoPEAlgo_t algo) {
+    auto result = RMSNormRoPEInfo::create(x_desc, w_desc, pos_desc, sin_desc, cos_desc, epsilon, algo);
+    CHECK_RESULT(result);
+    *desc_ptr = new Descriptor(nullptr, result.take(), 0, handle->device, handle->device_id);
+    return INFINI_STATUS_SUCCESS;
+}
+
+template <typename Tdata, typename Tweight, typename Tindex>
+infiniStatus_t rms_norm_rope(const RMSNormRoPEInfo *info,
+                             Tdata *x, const Tweight *w, const Tindex *pos_ids,
+                             const Tdata *sin_table, const Tdata *cos_table) {
+    const size_t num_tokens = info->num_tokens;
+    const size_t nhead = info->num_heads;
+    const size_t table_dim = info->table_dim;
+    const size_t dim = 2 * table_dim;
+    const ptrdiff_t total_blocks = static_cast<ptrdiff_t>(num_tokens * nhead);
+
+#pragma omp parallel for
+    for (ptrdiff_t block_idx = 0; block_idx < total_blocks; ++block_idx) {
+        const size_t tok = block_idx / nhead; // token index
+        const size_t h = block_idx % nhead;   // head index
+
+        Tdata *x_ptr = x + tok * info->x_stride_token + h * info->x_stride_head;
+
+        const size_t pos_id = size_t(pos_ids[tok * info->pos_stride]);
+        const Tdata *sin_ptr = sin_table + pos_id * table_dim;
+        const Tdata *cos_ptr = cos_table + pos_id * table_dim;
+
+        // [Reduce] sum of x^2 on the row
+        float ss = 0.f;
+        for (size_t k = 0; k < dim; k++) {
+            float v = utils::cast<float>(x_ptr[k]);
+            ss += v * v;
+        }
+
+        // 1 / (sqrt(sum/dim + eps))
+        float rms = 1.f / std::sqrt(ss / (float)(dim) + info->epsilon);
+
+        for (size_t i = 0; i < table_dim; i++) {
+            // Calculate positions based on algorithm
+            size_t pos0, pos1;
+            if (info->algo == infiniopRoPEAlgo_t::INFINIOP_ROPE_ALGO_GPT_J) {
+                // GPT-J style: interleaved pairs
+                pos0 = 2 * i;
+                pos1 = 2 * i + 1;
+            } else {
+                // GPT-NeoX style: first half and second half
+                pos0 = i;
+                pos1 = i + table_dim;
+            }
+
+            // Round normalized values to the storage dtype first
+            // (mimicking the two-kernel pipeline), then rotate in fp32
+            // with the same conversion and rounding order as the rope cpu kernel
+            float x0 = utils::cast<float>(utils::cast<Tdata>(utils::cast<float>(x_ptr[pos0]) * utils::cast<float>(w[pos0]) * rms));
+            float x1 = utils::cast<float>(utils::cast<Tdata>(utils::cast<float>(x_ptr[pos1]) * utils::cast<float>(w[pos1]) * rms));
+            float sin__ = utils::cast<float>(sin_ptr[i]);
+            float cos__ = utils::cast<float>(cos_ptr[i]);
+
+            x_ptr[pos0] = utils::cast<Tdata>(x0 * cos__ - x1 * sin__);
+            x_ptr[pos1] = utils::cast<Tdata>(x0 * sin__ + x1 * cos__);
+        }
+    }
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+#define CALCULATE(Tdata, Tweight, Tindex) \
+    rms_norm_rope(&_info, (Tdata *)x, (const Tweight *)w, (const Tindex *)pos_ids, (const Tdata *)sin_table, (const Tdata *)cos_table)
+
+#define DISPATCH_POS(Tdata, Tweight)                                   \
+    if (_info.pos_type == INFINI_DTYPE_I32) {                          \
+        return CALCULATE(Tdata, Tweight, int32_t);                     \
+    } else if (_info.pos_type == INFINI_DTYPE_I64) {                   \
+        return CALCULATE(Tdata, Tweight, int64_t);                     \
+    } else {                                                           \
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;                         \
+    }
+
+infiniStatus_t Descriptor::calculate(
+    void *workspace, size_t workspace_size,
+    void *x, const void *w,
+    const void *pos_ids,
+    const void *sin_table,
+    const void *cos_table,
+    void *stream) const {
+    if (_info.atype == INFINI_DTYPE_F16) {
+        if (_info.wtype == INFINI_DTYPE_F16) {
+            DISPATCH_POS(fp16_t, fp16_t);
+        } else if (_info.wtype == INFINI_DTYPE_BF16) {
+            DISPATCH_POS(fp16_t, bf16_t);
+        } else if (_info.wtype == INFINI_DTYPE_F32) {
+            DISPATCH_POS(fp16_t, float);
+        } else {
+            return INFINI_STATUS_BAD_TENSOR_DTYPE;
+        }
+    } else if (_info.atype == INFINI_DTYPE_BF16) {
+        if (_info.wtype == INFINI_DTYPE_BF16) {
+            DISPATCH_POS(bf16_t, bf16_t);
+        } else if (_info.wtype == INFINI_DTYPE_F16) {
+            DISPATCH_POS(bf16_t, fp16_t);
+        } else if (_info.wtype == INFINI_DTYPE_F32) {
+            DISPATCH_POS(bf16_t, float);
+        } else {
+            return INFINI_STATUS_BAD_TENSOR_DTYPE;
+        }
+    } else if (_info.atype == INFINI_DTYPE_F32) {
+        DISPATCH_POS(float, float);
+    } else {
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+#undef DISPATCH_POS
+#undef CALCULATE
+
+} // namespace op::rms_norm_rope::cpu

src/infiniop/ops/rms_norm_rope/cpu/rms_norm_rope_cpu.h

@@ -0,0 +1,7 @@
+#ifndef __RMS_NORM_ROPE_CPU_H__
+#define __RMS_NORM_ROPE_CPU_H__
+#include "../rms_norm_rope.h"
+
+DESCRIPTOR(cpu)
+
+#endif