Merge pull request #471 from InfiniTensor/fix/get-rope

fix: add basic get rope API and support mrope

Author: wooway777

csrc/layers/rotary_embedding/rotary_embedding.cpp

@@ -1,59 +1,115 @@
 #include "rotary_embedding.hpp"
 #include "../../config/model_config.hpp"
 #include "rotary_embedding_factory.hpp"
+#include <iomanip>
 #include <memory>
+#include <sstream>
 #include <string>
+#include <typeinfo>
+#include <unordered_map>
 
 namespace infinilm::layers::rotary_embedding {
 
+namespace {
 // Cache dictionary to avoid redundant allocations of RoPE instances.
 // thread_local ensures it is only visible within this compilation unit.
 thread_local std::unordered_map<std::string, std::shared_ptr<infinicore::nn::RoPE>> _ROPE_DICT;
 
-std::shared_ptr<infinicore::nn::RoPE>
-get_rope(const std::shared_ptr<infinilm::config::ModelConfig> &model_config,
-         const infinicore::Device &device) {
+std::string make_cache_key(size_t head_dim,
+                           size_t rotary_dim,
+                           size_t max_position_embeddings,
+                           double rope_theta,
+                           infinicore::nn::RoPE::Algo algo,
+                           const infinicore::DataType &dtype,
+                           const infinicore::Device &device,
+                           const std::shared_ptr<infinicore::nn::RopeScalingConfig> &scaling,
+                           const std::optional<std::vector<int>> &mrope_section,
+                           bool mrope_interleaved) {
+    std::ostringstream oss;
+    oss << "scaling_" << (scaling ? typeid(*scaling).name() : "none")
+        << "_head_" << head_dim
+        << "_rotary_" << rotary_dim
+        << "_max_" << max_position_embeddings
+        << "_theta_" << std::setprecision(17) << rope_theta
+        << "_algo_" << static_cast<int>(algo)
+        << "_dtype_" << static_cast<int>(dtype)
+        << "_dev" << device.toString();
 
-    // 1. Compute the actual rotary dimension
-    size_t rotary_dim = model_config->get_rotary_dim();
-    size_t head_dim = model_config->get_head_dim();
-
-    // 2. Resolve scaling config via the internal factory
-    auto scaling = make_scaling_config(model_config);
-
-    // 3. Cache key must include rotary_dim AND the actual scaling type
-    //    to avoid reusing the same RoPE instance across models with different settings
-    //    (Enhancement: dynamically determine scaling_type instead of hardcoding "default")
-    std::string scaling_type_str = "default";
-    if (scaling) {
-        // Assuming we can get the type string from the JSON for cache key generation,
-        // or ideally, ScalingConfig should have a virtual std::string type_name() const method.
-        // Here we read it from JSON for the cache key purpose only, keeping it decoupled from InfiniCore.
-        const auto &rope_scaling_json = model_config->get_config_json()["rope_scaling"];
-        if (rope_scaling_json.contains("type")) {
-            scaling_type_str = rope_scaling_json["type"].get<std::string>();
-        } else if (rope_scaling_json.contains("rope_type")) {
-            scaling_type_str = rope_scaling_json["rope_type"].get<std::string>();
+    if (mrope_section.has_value()) {
+        oss << "_mrope";
+        for (const auto section : mrope_section.value()) {
+            oss << '_' << section;
         }
+        oss << "_mrope_interleaved_" << (mrope_interleaved ? 1 : 0);
     }
+    return oss.str();
+}
+
+} // namespace
 
-    std::string cache_key = scaling_type_str + "_rope_dim_" + std::to_string(rotary_dim)
-                          + "_dev_" + device.toString();
+std::shared_ptr<infinicore::nn::RoPE>
+get_rope(size_t head_dim,
+         size_t rotary_dim,
+         size_t max_position_embeddings,
+         double rope_theta,
+         infinicore::nn::RoPE::Algo algo,
+         const infinicore::DataType &dtype,
+         const infinicore::Device &device,
+         std::shared_ptr<infinicore::nn::RopeScalingConfig> scaling,
+         std::optional<std::vector<int>> mrope_section,
+         bool mrope_interleaved) {
+    const auto cache_key = make_cache_key(head_dim,
+                                          rotary_dim,
+                                          max_position_embeddings,
+                                          rope_theta,
+                                          algo,
+                                          dtype,
+                                          device,
+                                          scaling,
+                                          mrope_section,
+                                          mrope_interleaved);
     auto it = _ROPE_DICT.find(cache_key);
     if (it != _ROPE_DICT.end()) {
         return it->second;
     }
 
+    auto rope = std::make_shared<infinicore::nn::RoPE>(head_dim,
+                                                       rotary_dim,
+                                                       max_position_embeddings,
+                                                       rope_theta,
+                                                       algo,
+                                                       dtype,
+                                                       device,
+                                                       scaling,
+                                                       mrope_section,
+                                                       mrope_interleaved);
+    _ROPE_DICT.emplace(cache_key, rope);
+    return rope;
+}
+
+std::shared_ptr<infinicore::nn::RoPE>
+get_rope(const std::shared_ptr<infinilm::config::ModelConfig> &model_config,
+         const infinicore::Device &device) {
+
+    size_t rotary_dim = model_config->get_rotary_dim();
+    size_t head_dim = model_config->get_head_dim();
+    auto scaling = make_scaling_config(model_config);
+
     const auto &dtype = model_config->get_dtype();
     size_t max_position_embeddings = model_config->get<size_t>("max_position_embeddings");
     double rope_theta = model_config->get<double>("rope_theta");
-
     infinicore::nn::RoPE::Algo algo = model_config->get_rope_algo();
-    auto rope = std::make_shared<infinicore::nn::RoPE>(head_dim, rotary_dim, max_position_embeddings, rope_theta,
-                                                       algo, dtype, device, scaling);
 
-    _ROPE_DICT.emplace(cache_key, rope);
-    return rope;
+    return get_rope(head_dim,
+                    rotary_dim,
+                    max_position_embeddings,
+                    rope_theta,
+                    algo,
+                    dtype,
+                    device,
+                    scaling,
+                    std::nullopt,
+                    false);
 }
 
 } // namespace infinilm::layers::rotary_embedding

csrc/layers/rotary_embedding/rotary_embedding.hpp

@@ -2,19 +2,35 @@
 
 #include "infinicore/nn/rope.hpp"
 #include <memory>
+#include <optional>
+#include <string>
+#include <vector>
 
 namespace infinilm::config {
 class ModelConfig; // Forward declaration
 }
 
 namespace infinilm::layers::rotary_embedding {
 
+std::shared_ptr<infinicore::nn::RoPE>
+get_rope(size_t head_dim,
+         size_t rotary_dim,
+         size_t max_position_embeddings,
+         double rope_theta,
+         infinicore::nn::RoPE::Algo algo,
+         const infinicore::DataType &dtype,
+         const infinicore::Device &device,
+         std::shared_ptr<infinicore::nn::RopeScalingConfig> scaling = nullptr,
+         std::optional<std::vector<int>> mrope_section = std::nullopt,
+         bool mrope_interleaved = false);
+
 /**
  * @brief Public API to assemble and construct a complete RoPE module.
  *
  * @param model_config Model configuration.
  * @param device Device to create the cache on.
- * @param algo RoPE algorithm type (default: Algo::GPT_NEOX).
+ * @todo Not Supported Yet: reading mrope params from config
+ *
  */
 std::shared_ptr<infinicore::nn::RoPE>
 get_rope(const std::shared_ptr<infinilm::config::ModelConfig> &model_config,