refactor(ops): route basic Llama through canonical InfiniOps

include/infinicore/nn/rope.hpp

@@ -94,9 +94,10 @@ class RoPE : public Module {
     std::string extra_repr() const;
 
 protected:
-    // Buffers (sin and cos cache tables) - not exposed in state_dict
+    // RoPE cache tables are not exposed in state_dict.
     INFINICORE_NN_BUFFER(sin_cache);
     INFINICORE_NN_BUFFER(cos_cache);
+    INFINICORE_NN_BUFFER(cos_sin_cache);
 
 private:
     void initialize_cache();

include/infinicore/ops.hpp

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

include/infinicore/ops/rotary_embedding.hpp

@@ -0,0 +1,23 @@
+#pragma once
+
+#include "../graph/graph.hpp"
+#include "../tensor.hpp"
+#include "common/op.hpp"
+
+#include <cstdint>
+#include <optional>
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(RotaryEmbedding, const Tensor &, Tensor, std::optional<Tensor>, const Tensor &, int64_t, bool, int64_t, bool);
+
+void rotary_embedding_(const Tensor &positions,
+                       Tensor query,
+                       std::optional<Tensor> key,
+                       const Tensor &cos_sin_cache,
+                       int64_t head_size,
+                       bool is_neox,
+                       int64_t rope_dim_offset = 0,
+                       bool inverse = false);
+
+} // namespace infinicore::op

src/infinicore/nn/rope.cc

@@ -3,6 +3,7 @@
 #include "../utils.hpp"
 #include "infinicore/ops/mrope.hpp"
 #include "infinicore/ops/rope.hpp"
+#include "infinicore/ops/rotary_embedding.hpp"
 #include <algorithm>
 #include <cassert>
 #include <cmath>
@@ -12,6 +13,24 @@
 #include <vector>
 
 namespace infinicore::nn {
+namespace {
+
+template <typename T>
+std::vector<T> combine_cos_sin_cache(const std::vector<T> &cos,
+                                     const std::vector<T> &sin,
+                                     size_t max_seq_len,
+                                     size_t cache_dim) {
+    std::vector<T> combined(max_seq_len * cache_dim * 2);
+    for (size_t pos = 0; pos < max_seq_len; ++pos) {
+        std::copy_n(cos.begin() + pos * cache_dim, cache_dim,
+                    combined.begin() + pos * cache_dim * 2);
+        std::copy_n(sin.begin() + pos * cache_dim, cache_dim,
+                    combined.begin() + (pos * 2 + 1) * cache_dim);
+    }
+    return combined;
+}
+
+} // namespace
 
 RoPE::RoPE(size_t head_dim,
            size_t rotary_dim,
@@ -58,6 +77,12 @@ void RoPE::initialize_cache() {
     INFINICORE_NN_BUFFER_INIT(sin_cache, ({max_seq_len_, cache_dim}, dtype_, device_));
     INFINICORE_NN_BUFFER_INIT(cos_cache, ({max_seq_len_, cache_dim}, dtype_, device_));
 
+#ifdef ENABLE_INFINIOPS_API
+    if (device_.getType() == Device::Type::NVIDIA && !mrope_section_) {
+        INFINICORE_NN_BUFFER_INIT(cos_sin_cache, ({max_seq_len_, rotary_dim_}, dtype_, device_));
+    }
+#endif
+
     // Pre-compute sin and cos values
     // Frequency generation always uses GPT-J style (theta^(-2j/rotary_dim)).
     // The rotation algorithm (algo_) controls how dimensions are paired in the kernel.
@@ -95,6 +120,11 @@ void RoPE::initialize_cache() {
         auto cos_f32_cpu = Tensor::from_blob(cos_data.data(), {max_seq_len_, cache_dim}, DataType::F32, cpu_device);
         sin_cache_->copy_from(sin_f32_cpu);
         cos_cache_->copy_from(cos_f32_cpu);
+        if (cos_sin_cache_) {
+            auto combined = combine_cos_sin_cache(cos_data, sin_data, max_seq_len_, cache_dim);
+            auto combined_cpu = Tensor::from_blob(combined.data(), {max_seq_len_, rotary_dim_}, DataType::F32, cpu_device);
+            cos_sin_cache_->copy_from(combined_cpu);
+        }
     } else if (dtype_ == DataType::BF16) {
         // Convert F32 to BF16 using the same conversion as Python's ml_dtypes.bfloat16
         // This uses round-to-nearest-even (matching _f32_to_bf16 implementation)
@@ -112,6 +142,11 @@ void RoPE::initialize_cache() {
         // copy_from handles cross-device copying to target device
         sin_cache_->copy_from(sin_bf16_cpu);
         cos_cache_->copy_from(cos_bf16_cpu);
+        if (cos_sin_cache_) {
+            auto combined = combine_cos_sin_cache(cos_bf16_data, sin_bf16_data, max_seq_len_, cache_dim);
+            auto combined_cpu = Tensor::from_blob(combined.data(), {max_seq_len_, rotary_dim_}, DataType::BF16, cpu_device);
+            cos_sin_cache_->copy_from(combined_cpu);
+        }
     } else if (dtype_ == DataType::F16) {
         // Convert F32 to F16
         std::vector<fp16_t> sin_f16_data(max_seq_len_ * cache_dim);
@@ -127,6 +162,11 @@ void RoPE::initialize_cache() {
 
         sin_cache_->copy_from(sin_f16_cpu);
         cos_cache_->copy_from(cos_f16_cpu);
+        if (cos_sin_cache_) {
+            auto combined = combine_cos_sin_cache(cos_f16_data, sin_f16_data, max_seq_len_, cache_dim);
+            auto combined_cpu = Tensor::from_blob(combined.data(), {max_seq_len_, rotary_dim_}, DataType::F16, cpu_device);
+            cos_sin_cache_->copy_from(combined_cpu);
+        }
     } else {
         throw std::runtime_error(
             "RoPE cache dtype conversion not yet supported for dtype: "
@@ -148,6 +188,21 @@ Tensor RoPE::forward(const Tensor &x, const Tensor &pos, bool in_place) const {
         }
     }
 
+#ifdef ENABLE_INFINIOPS_API
+    if (cos_sin_cache_) {
+        if (!in_place) {
+            y->copy_from(x);
+        }
+        op::rotary_embedding_(pos,
+                              y,
+                              std::nullopt,
+                              cos_sin_cache_,
+                              static_cast<int64_t>(head_dim_),
+                              algo_ == Algo::GPT_NEOX);
+        return y;
+    }
+#endif
+
     size_t ndim = x->ndim();
     op::rope_(y->narrow({{ndim - 1, 0, rotary_dim_}}),
               x->narrow({{ndim - 1, 0, rotary_dim_}}),

src/infinicore/ops/add_rms_norm/add_rms_norm_infiniops.cc

@@ -3,7 +3,8 @@
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/add_rms_norm.h"
+#include "base/copy.h"
+#include "base/fused_add_rms_norm.h"
 
 #include <optional>
 
@@ -16,6 +17,7 @@ struct PlannedMeta {
     TensorMeta out, residual, a, b, weight;
     graph::GraphTensor out_tensor, residual_tensor, a_tensor, b_tensor, weight_tensor;
     float epsilon;
+    bool copy_input, copy_residual;
 };
 
 } // namespace
@@ -35,7 +37,9 @@ void *plan(Tensor out, Tensor residual, const Tensor &a, const Tensor &b, const
         graph::GraphTensor(a),
         graph::GraphTensor(b),
         graph::GraphTensor(weight),
-        epsilon};
+        epsilon,
+        out->data() != a->data(),
+        residual->data() != b->data()};
 }
 
 void run(void *planned_meta) {
@@ -45,15 +49,23 @@ void run(void *planned_meta) {
     handle.set_stream(context::getStream());
     infini::ops::Config config;
 
-    infini::ops::AddRmsNorm::Call(
+    auto out = planned->out.tensor(planned->out_tensor);
+    auto residual = planned->residual.tensor(planned->residual_tensor);
+
+    if (planned->copy_input) {
+        infini::ops::Copy::Call(handle, config, planned->a.tensor(planned->a_tensor), false, out);
+    }
+    if (planned->copy_residual) {
+        infini::ops::Copy::Call(handle, config, planned->b.tensor(planned->b_tensor), false, residual);
+    }
+
+    infini::ops::FusedAddRmsNorm::Call(
         handle,
         config,
-        planned->a.tensor(planned->a_tensor),
-        planned->b.tensor(planned->b_tensor),
-        planned->weight.tensor(planned->weight_tensor),
-        std::optional<float>{planned->epsilon},
-        planned->out.tensor(planned->out_tensor),
-        planned->residual.tensor(planned->residual_tensor));
+        out,
+        residual,
+        std::optional<infini::ops::Tensor>{planned->weight.tensor(planned->weight_tensor)},
+        planned->epsilon);
 }
 
 void cleanup(void **planned_meta_ptr) {