issue/545 nn::module::Rope and nn::module::Swiglu

Signed-off-by: Ceng23333 <441651826@qq.com>

Author: Ceng23333

include/infinicore/nn/rope.hpp

@@ -0,0 +1,77 @@
+#pragma once
+
+#include "module.hpp"
+#include "../context/context.hpp"
+#include "../tensor.hpp"
+#include <infiniop.h>
+#include <memory>
+
+namespace infinicore::nn {
+
+
+class RoPE : public Module {
+public:
+    /**
+     * @brief Construct a RoPE layer
+     *
+     * @param head_dim Dimension of each attention head (must be even)
+     * @param max_seq_len Maximum sequence length for pre-computed cache
+     * @param theta Base frequency for rotary embeddings (default: 10000.0)
+     * @param algo RoPE algorithm type (default: INFINIOP_ROPE_ALGO_GPT_J)
+     * @param dtype Data type for sin/cos cache (default: DataType::F32)
+     * @param device Device to create the cache on
+     */
+    RoPE(size_t head_dim,
+         size_t max_seq_len,
+         double theta = 10000.0,
+         infiniopRoPEAlgo_t algo = INFINIOP_ROPE_ALGO_GPT_J,
+         const DataType &dtype = DataType::F32,
+         const Device &device = Device());
+
+    /**
+     * @brief Forward pass: apply RoPE to a tensor
+     *
+     * @param x Input tensor of shape (..., head_dim) where ... is any number of dimensions
+     * @param pos Position IDs tensor of shape (*,) typically [seq_len] or [batch, seq_len]
+     * @return Rotated tensor with same shape as input
+     *
+     * Applies rotary position embeddings to the input tensor.
+     * For attention mechanisms, call this method separately for query and key tensors.
+     *
+     * Common input shapes:
+     *   - [batch, num_heads, seq_len, head_dim]
+     *   - [batch, seq_len, num_heads, head_dim]
+     *   - [seq_len, head_dim]
+     */
+    Tensor forward(const Tensor &x, const Tensor &pos) const;
+
+    // Module information
+    size_t head_dim() const { return head_dim_; }
+    size_t max_seq_len() const { return max_seq_len_; }
+    double theta() const { return theta_; }
+    infiniopRoPEAlgo_t algo() const { return algo_; }
+    DataType dtype() const { return dtype_; }
+
+    // String representation
+    std::string extra_repr() const;
+
+    // Accessors for parameters
+    Tensor sin_cache() const { return sin_cache_; }
+    Tensor cos_cache() const { return cos_cache_; }
+
+protected:
+    // Parameters (sin and cos cache tables)
+    INFINICORE_NN_PARAMETER(sin_cache);
+    INFINICORE_NN_PARAMETER(cos_cache);
+
+private:
+    void initialize_cache();
+
+    size_t head_dim_;      // Dimension of each attention head
+    size_t max_seq_len_;   // Maximum sequence length
+    double theta_;         // Base frequency for rotary embeddings
+    infiniopRoPEAlgo_t algo_;  // RoPE algorithm type
+    DataType dtype_;       // Data type for cache tables
+};
+
+} // namespace infinicore::nn

include/infinicore/nn/swiglu.hpp

@@ -0,0 +1,35 @@
+#pragma once
+
+#include "module.hpp"
+#include "../ops.hpp"
+
+namespace infinicore::nn {
+
+class SwiGLU : public Module {
+public:
+    /**
+     * @brief Construct a SwiGLU module
+     *
+     * SwiGLU is a stateless activation function, so no parameters are needed.
+     */
+    SwiGLU() = default;
+
+    /**
+     * @brief Forward pass: apply SwiGLU activation
+     *
+     * @param up The "up" projection tensor
+     * @param gate The "gate" projection tensor
+     * @return Result tensor: up * gate * sigmoid(gate)
+     *
+     * Both input tensors must have the same shape and dtype.
+     * Common usage:
+     *   - Input: up from linear_up layer, gate from linear_gate layer
+     *   - Shapes: typically [batch, seq_len, hidden_size] or [batch, hidden_size]
+     */
+    Tensor forward(const Tensor &up, const Tensor &gate) const;
+
+    // String representation
+    std::string extra_repr() const;
+};
+
+} // namespace infinicore::nn

include/infinicore/ops.hpp

@@ -7,5 +7,6 @@
 #include "ops/ones.hpp"
 #include "ops/rearrange.hpp"
 #include "ops/rms_norm.hpp"
+#include "ops/rope.hpp"
 #include "ops/silu.hpp"
 #include "ops/swiglu.hpp"

include/infinicore/ops/rope.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "../tensor.hpp"
+#include "common/op.hpp"
+#include <infiniop.h>
+
+namespace infinicore::op {
+class RoPE {
+public:
+    using schema = void (*)(Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, infiniopRoPEAlgo_t);
+    static void execute(Tensor x_out, const Tensor &x, const Tensor &pos, const Tensor &sin_cache, const Tensor &cos_cache, infiniopRoPEAlgo_t algo);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+void rope_(Tensor x_out, const Tensor &x, const Tensor &pos, const Tensor &sin_cache, const Tensor &cos_cache, infiniopRoPEAlgo_t algo);
+Tensor rope(const Tensor &x, const Tensor &pos, const Tensor &sin_cache, const Tensor &cos_cache, infiniopRoPEAlgo_t algo);
+} // namespace infinicore::op

src/infinicore-test/main.cc

@@ -141,12 +141,8 @@ ParsedArgs parseArgs(int argc, char *argv[]) {
 
 int main(int argc, char *argv[]) {
     try {
-        // Initialize spdlog for debugging
-        spdlog::set_level(spdlog::level::debug);
-        spdlog::info("Starting InfiniCore Memory Management Test Suite");
-
         ParsedArgs args = parseArgs(argc, argv);
-        spdlog::debug("Arguments parsed successfully");
+        spdlog::info("Arguments parsed successfully");
 
         std::cout << "==============================================\n"
                   << "InfiniCore Memory Management Test Suite\n"
@@ -156,31 +152,25 @@ int main(int argc, char *argv[]) {
                   << "Iterations: " << args.iterations << "\n"
                   << "==============================================" << std::endl;
 
-        spdlog::debug("About to initialize InfiniCore context");
+        spdlog::info("About to initialize InfiniCore context");
         // Initialize InfiniCore context
         infinicore::context::setDevice(infinicore::Device(static_cast<infinicore::Device::Type>(args.device_type), 0));
-        spdlog::debug("InfiniCore context initialized successfully");
+        spdlog::info("InfiniCore context initialized successfully");
 
-        spdlog::debug("Creating test runner");
+        spdlog::info("Creating test runner");
         // Create test runner
         infinicore::test::InfiniCoreTestRunner runner;
-        spdlog::debug("Test runner created successfully");
+        spdlog::info("Test runner created successfully");
 
         // Add tests based on arguments
         if (args.run_basic) {
-            spdlog::debug("Adding BasicMemoryTest");
             runner.addTest(std::make_unique<infinicore::test::BasicMemoryTest>());
-            spdlog::debug("BasicMemoryTest added successfully");
 
-            spdlog::debug("Adding TensorDestructorTest");
             runner.addTest(std::make_unique<infinicore::test::TensorDestructorTest>());
-            spdlog::debug("TensorDestructorTest added successfully");
         }
 
         if (args.run_module) {
-            spdlog::debug("Adding NNModuleTest");
             runner.addTest(std::make_unique<infinicore::test::NNModuleTest>());
-            spdlog::debug("NNModuleTest added successfully");
         }
 
         if (args.run_concurrency) {
@@ -203,10 +193,10 @@ int main(int argc, char *argv[]) {
             runner.addTest(std::make_unique<infinicore::test::StressTest>());
         }
 
-        spdlog::debug("About to run all tests");
+        spdlog::info("About to run all tests");
         // Run all tests
         auto results = runner.runAllTests();
-        spdlog::debug("All tests completed");
+        spdlog::info("All tests completed");
 
         // Count results and collect failed tests
         size_t passed = 0, failed = 0;