Pr 521

csrc/engine/compiler/paged_compiler.cpp

@@ -62,6 +62,14 @@ void PagedCompiler::compile() {
         auto &forward_context = infinilm::global_state::get_forward_context();
         const bool has_mamba_state = has_mamba_cache(forward_context);
 
+        const auto &model_config = model_->get_model_config();
+        const size_t position_id_axes = model_config == nullptr
+                                          ? 1
+                                          : model_config->get_or<size_t>("position_id_axes", 1);
+        if (position_id_axes == 0) {
+            throw std::runtime_error("PagedCompiler: position_id_axes must be positive");
+        }
+
         size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end());
         compiled_map_decode_.clear();
         block_tables_holder_ = infinicore::Tensor::empty(
@@ -71,7 +79,11 @@ void PagedCompiler::compile() {
         auto make_decode_input = [&](size_t b) {
             InfinilmModel::Input input;
             input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice());
-            input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
+            input.position_ids = infinicore::Tensor::empty(
+                position_id_axes > 1
+                    ? std::vector<size_t>{position_id_axes, b}
+                    : std::vector<size_t>{b},
+                infinicore::DataType::I64, infinicore::context::getDevice());
             input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice());
             set_zeros(input.input_ids.value());
             set_zeros(input.position_ids.value());

csrc/models/infinilm_model.hpp

@@ -70,6 +70,9 @@ class InfinilmModel : public infinicore::nn::Module {
     virtual const cache::CacheConfig *get_cache_config() const {
         return cache_config_.get();
     }
+    const std::shared_ptr<infinilm::config::ModelConfig> &get_model_config() const {
+        return model_config_;
+    }
 
     void process_weights_after_loading();
     void reset_runtime_state() const;

csrc/models/qwen3_5/qwen3_5_attention.cpp

@@ -32,12 +32,15 @@ Qwen35Attention::Qwen35Attention(std::shared_ptr<infinilm::config::ModelConfig>
     const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info();
     int tp_rank = infinilm::global_state::get_tensor_model_parallel_rank();
     int tp_size = infinilm::global_state::get_tensor_model_parallel_world_size();
-    if ((total_num_kv_heads < tp_size) || (0 != (total_num_kv_heads % tp_size))) {
+    if (total_num_kv_heads >= static_cast<size_t>(tp_size)
+        && total_num_kv_heads % tp_size != 0) {
         throw std::runtime_error("infinilm::models::qwen3_5::Qwen35Attention: num_key_value_heads must be divisible by tp_size");
     }
 
     num_attention_heads_ = total_num_heads / tp_size;
-    num_key_value_heads_ = total_num_kv_heads / tp_size;
+    num_key_value_heads_ = total_num_kv_heads < static_cast<size_t>(tp_size)
+                             ? 1
+                             : total_num_kv_heads / tp_size;
 
     auto quantization_method = model_config->get_quantization_method();
     auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); };

csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp

@@ -1,32 +1,33 @@
 #include "qwen3_5_for_causal_lm.hpp"
 
-#include "../../global_state/global_state.hpp"
 #include "../models_registry.hpp"
-#include "../qwen3_next/qwen3_next_for_causal_lm.hpp"
 #include <stdexcept>
 #include <string>
 #include <vector>
 
 namespace infinilm::models::qwen3_5 {
 
-Qwen35ForCausalLM::Qwen35ForCausalLM(std::shared_ptr<infinilm::config::ModelConfig> model_config,
-                                     const infinicore::Device &device) {
+Qwen35ForCausalLM::Qwen35ForCausalLM(
+    std::shared_ptr<infinilm::config::ModelConfig> model_config,
+    const infinicore::Device &device) {
     model_config_ = model_config;
-    size_t hidden_size = model_config->get<size_t>("hidden_size");
-    size_t vocab_size = model_config->get<size_t>("vocab_size");
-    const auto &dtype{model_config->get_dtype()};
+    const size_t hidden_size = model_config->get<size_t>("hidden_size");
+    const size_t vocab_size = model_config->get<size_t>("vocab_size");
+    const auto &dtype = model_config->get_dtype();
 
     INFINICORE_NN_MODULE_INIT(model, model_config, device);
-    INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device);
+    INFINICORE_NN_MODULE_INIT(
+        lm_head, hidden_size, vocab_size, false, dtype, device);
 }
 
-infinilm::InfinilmModel::Output Qwen35ForCausalLM::forward(const infinilm::InfinilmModel::Input &input) const {
+InfinilmModel::Output Qwen35ForCausalLM::forward(
+    const InfinilmModel::Input &input) const {
     auto hidden_states = model_->forward(input);
-    auto logits = lm_head_->forward(hidden_states);
-    return {logits};
+    return {lm_head_->forward(hidden_states)};
 }
 
-void Qwen35ForCausalLM::reset_cache(const cache::CacheConfig *cache_config) {
+void Qwen35ForCausalLM::reset_cache(
+    const cache::CacheConfig *cache_config) {
     if (cache_config == nullptr) {
         cache_config_.reset();
     } else {
@@ -35,12 +36,7 @@ void Qwen35ForCausalLM::reset_cache(const cache::CacheConfig *cache_config) {
     model_->reset_cache(cache_config);
 }
 
-std::shared_ptr<infinilm::config::ModelConfig> create_qwen3_5_model_config(std::shared_ptr<infinilm::config::ModelConfig> model_config) {
-    const std::string model_type = model_config->get<std::string>("model_type");
-    if ("qwen3_5" != model_type) {
-        throw std::runtime_error("infinilm::models::qwen3_5::create_qwen3_next_model_config: model_type is not qwen3_5");
-    }
-
+std::shared_ptr<infinilm::config::ModelConfig> prepare_qwen3_5_model_config(std::shared_ptr<infinilm::config::ModelConfig> model_config) {
     nlohmann::json &config_json = model_config->get_config_json();
     if (config_json.contains("text_config") && config_json["text_config"].is_object()) {
         const nlohmann::json &text_config_json = config_json["text_config"];
@@ -53,32 +49,56 @@ std::shared_ptr<infinilm::config::ModelConfig> create_qwen3_5_model_config(std::
             config_json["dtype"] = config_json["torch_dtype"];
         }
     }
+    if (!config_json.contains("position_id_axes")) {
+        size_t position_id_axes = 1;
+        if (config_json.contains("rope_parameters")
+            && config_json["rope_parameters"].is_object()) {
+            const auto &rope_parameters = config_json["rope_parameters"];
+            if (rope_parameters.contains("mrope_section")
+                && rope_parameters["mrope_section"].is_array()
+                && !rope_parameters["mrope_section"].empty()) {
+                position_id_axes = rope_parameters["mrope_section"].size();
+            }
+        }
+        config_json["position_id_axes"] = position_id_axes;
+    }
     if (!config_json.contains("rope_theta") && config_json.contains("rope_parameters") && config_json["rope_parameters"].is_object() && config_json["rope_parameters"].contains("rope_theta")) {
-        // TODO: This is only a temporary loader shim. Qwen3.6 uses mRoPE,
-        // which needs proper support in InfiniCore instead of treating it as
-        // plain RoPE through a top-level rope_theta.
+        // Normalize the nested HuggingFace field for the Qwen3.5 attention module.
         config_json["rope_theta"] = config_json["rope_parameters"]["rope_theta"];
     }
     if (!config_json.contains("partial_rotary_factor") && config_json.contains("rope_parameters") && config_json["rope_parameters"].is_object() && config_json["rope_parameters"].contains("partial_rotary_factor")) {
         config_json["partial_rotary_factor"] = config_json["rope_parameters"]["partial_rotary_factor"];
     }
     if (!config_json.contains("layer_types")) {
-        size_t full_attention_interval = model_config->get<size_t>("full_attention_interval");
-        size_t num_hidden_layers = model_config->get<size_t>("num_hidden_layers");
+        const size_t full_attention_interval = model_config->get<size_t>("full_attention_interval");
+        if (full_attention_interval == 0) {
+            throw std::runtime_error("Qwen3.5 full_attention_interval must be positive");
+        }
+        const size_t num_hidden_layers = model_config->get<size_t>("num_hidden_layers");
         std::vector<std::string> layer_types;
         layer_types.reserve(num_hidden_layers);
-        for (size_t i = 0; i < num_hidden_layers; i++) {
-            layer_types.push_back(bool((i + 1) % full_attention_interval) ? "linear_attention" : "full_attention");
+        for (size_t i = 0; i < num_hidden_layers; ++i) {
+            layer_types.push_back(
+                (i + 1) % full_attention_interval == 0
+                    ? "full_attention"
+                    : "linear_attention");
         }
-        config_json["layer_types"] = layer_types;
+        config_json["layer_types"] = std::move(layer_types);
     }
-
     if (!config_json.contains("attention_bias")) {
         config_json["attention_bias"] = false;
     }
     return model_config;
 }
 
+std::shared_ptr<infinilm::config::ModelConfig> create_qwen3_5_model_config(std::shared_ptr<infinilm::config::ModelConfig> model_config) {
+    const std::string model_type = model_config->get<std::string>("model_type");
+    if ("qwen3_5" != model_type) {
+        throw std::runtime_error("infinilm::models::qwen3_5::create_qwen3_5_model_config: model_type is not qwen3_5");
+    }
+    return prepare_qwen3_5_model_config(model_config);
+}
+
 } // namespace infinilm::models::qwen3_5
 
 namespace {

csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp

@@ -20,6 +20,9 @@ class Qwen35ForCausalLM : public InfinilmModel {
     INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head);
 };
 
+std::shared_ptr<infinilm::config::ModelConfig> prepare_qwen3_5_model_config(
+    std::shared_ptr<infinilm::config::ModelConfig> model_config);
+
 std::shared_ptr<infinilm::config::ModelConfig> create_qwen3_5_model_config(std::shared_ptr<infinilm::config::ModelConfig> model_config);
 
 } // namespace infinilm::models::qwen3_5

csrc/models/qwen3_5/qwen3_5_model.cpp

@@ -31,20 +31,19 @@ std::vector<int32_t> tensor_to_i32_vector(const infinicore::Tensor &tensor) {
 
 } // namespace
 
-Qwen35Model::Qwen35Model(std::shared_ptr<infinilm::config::ModelConfig> model_config,
-                         const infinicore::Device &device)
+Qwen35ModelBase::Qwen35ModelBase(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                                 const infinicore::Device &device)
     : model_config_(model_config) {
     const auto &dtype{model_config->get_dtype()};
     nlohmann::json &config_json = model_config->get_config_json();
 
     if (config_json.contains("vision_config") && !config_json["vision_config"].is_null()) {
         INFINICORE_NN_MODULE_INIT(visual, config_json["vision_config"], dtype, device);
     }
-    INFINICORE_NN_MODULE_INIT(language_model, model_config, device);
 }
 
-void Qwen35Model::replace_image_embeddings(infinicore::Tensor &inputs_embeds,
-                                           const InfinilmModel::Input &input) const {
+void Qwen35ModelBase::replace_image_embeddings(infinicore::Tensor &inputs_embeds,
+                                               const InfinilmModel::Input &input) const {
     if (!input.pixel_values.has_value() || input.pixel_values->empty()) {
         return;
     }
@@ -107,28 +106,19 @@ void Qwen35Model::replace_image_embeddings(infinicore::Tensor &inputs_embeds,
     }
 }
 
-infinicore::Tensor Qwen35Model::forward(const InfinilmModel::Input &input) const {
-    if (input.pixel_values.has_value() && !input.pixel_values->empty()) {
-        auto inputs_embeds = language_model_->embed_tokens(input.input_ids.value());
-        replace_image_embeddings(inputs_embeds, input);
-        return language_model_->forward_embeds(inputs_embeds, input.position_ids.value());
-    }
-    return language_model_->forward(input);
-}
-
-void Qwen35Model::reset_cache(const cache::CacheConfig *cache_config) {
-    if (nullptr == cache_config) {
-        return;
-    }
-
+void Qwen35ModelBase::reset_cache(const cache::CacheConfig *cache_config) {
     auto &forward_context = infinilm::global_state::get_forward_context();
     forward_context.kv_cache_vec.clear();
     forward_context.conv_state_vec.clear();
     forward_context.ssm_state_vec.clear();
+    if (nullptr == cache_config) {
+        return;
+    }
 
     const backends::AttentionBackend attention_backend = infinilm::global_state::get_infinilm_config().attention_backend;
 
-    auto cache_vectors = infinilm::models::qwen3_next::qwen3_next_allocate_cache_tensors(cache_config, model_config_, attention_backend);
+    auto cache_vectors = infinilm::models::qwen3_next::qwen3_next_allocate_cache_tensors(
+        cache_config, model_config_, attention_backend);
     forward_context.kv_cache_vec = std::move(cache_vectors.kv_cache_tensors);
     forward_context.conv_state_vec = std::move(cache_vectors.conv_state_tensors);
     forward_context.ssm_state_vec = std::move(cache_vectors.ssm_state_tensors);

csrc/models/qwen3_5/qwen3_5_model.hpp

@@ -10,24 +10,46 @@ namespace infinilm::models::qwen3_5 {
 
 using Qwen35LanguageModel = infinilm::layers::causal_lm_templates::TextModel<Qwen35DecoderLayer>;
 
-class Qwen35Model : public infinicore::nn::Module {
+class Qwen35ModelBase : public infinicore::nn::Module {
 public:
-    Qwen35Model(std::shared_ptr<infinilm::config::ModelConfig> model_config,
-                const infinicore::Device &device);
-
-    infinicore::Tensor forward(const InfinilmModel::Input &input) const;
+    Qwen35ModelBase(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                    const infinicore::Device &device);
 
     void reset_cache(const cache::CacheConfig *cache_config);
 
-private:
+protected:
     void replace_image_embeddings(infinicore::Tensor &inputs_embeds,
                                   const infinilm::InfinilmModel::Input &input) const;
 
-protected:
     INFINICORE_NN_MODULE(Qwen35VisionModel, visual);
-    INFINICORE_NN_MODULE(Qwen35LanguageModel, language_model);
-
     std::shared_ptr<infinilm::config::ModelConfig> model_config_;
 };
 
+template <typename LanguageModel>
+class Qwen35ModelTemplate : public Qwen35ModelBase {
+public:
+    Qwen35ModelTemplate(
+        std::shared_ptr<infinilm::config::ModelConfig> model_config,
+        const infinicore::Device &device)
+        : Qwen35ModelBase(model_config, device) {
+        language_model_ = this->register_module<LanguageModel>(
+            "language_model", model_config, device);
+    }
+
+    infinicore::Tensor forward(const InfinilmModel::Input &input) const {
+        if (input.pixel_values.has_value() && !input.pixel_values->empty()) {
+            auto inputs_embeds = language_model_->embed_tokens(input.input_ids.value());
+            replace_image_embeddings(inputs_embeds, input);
+            return language_model_->forward_embeds(
+                inputs_embeds, input.position_ids.value());
+        }
+        return language_model_->forward(input);
+    }
+
+protected:
+    INFINICORE_NN_MODULE(LanguageModel, language_model);
+};
+
+using Qwen35Model = Qwen35ModelTemplate<Qwen35LanguageModel>;
+
 } // namespace infinilm::models::qwen3_5

csrc/models/qwen3_5_moe/qwen3_5_moe_decoder_layer.cpp

@@ -0,0 +1,71 @@
+#include "qwen3_5_moe_decoder_layer.hpp"
+
+#include "infinicore/ops.hpp"
+
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace infinilm::models::qwen3_5_moe {
+
+Qwen35MoeDecoderLayer::Qwen35MoeDecoderLayer(
+    std::shared_ptr<infinilm::config::ModelConfig> model_config,
+    size_t layer_idx,
+    const infinicore::Device &device)
+    : layer_idx_(layer_idx) {
+    const auto &dtype = model_config->get_dtype();
+    const size_t hidden_size = model_config->get<size_t>("hidden_size");
+    const double rms_norm_eps = model_config->get<double>("rms_norm_eps");
+
+    INFINICORE_NN_MODULE_INIT(input_layernorm, hidden_size, rms_norm_eps, dtype, device);
+    INFINICORE_NN_MODULE_INIT(post_attention_layernorm, hidden_size, rms_norm_eps, dtype, device);
+    INFINICORE_NN_MODULE_INIT(mlp, model_config, layer_idx, device);
+
+    const auto layer_types = model_config->get<std::vector<std::string>>("layer_types");
+    layer_type_ = layer_types.at(layer_idx);
+    if ("linear_attention" == layer_type_) {
+        INFINICORE_NN_MODULE_INIT(linear_attn, model_config, layer_idx, device);
+    } else if ("full_attention" == layer_type_) {
+        INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device);
+    } else {
+        throw std::runtime_error(
+            "Qwen35MoeDecoderLayer: unsupported layer_type '" + layer_type_
+            + "' for layer " + std::to_string(layer_idx));
+    }
+}
+
+std::tuple<infinicore::Tensor, infinicore::Tensor> Qwen35MoeDecoderLayer::forward(
+    const infinicore::Tensor &positions,
+    infinicore::Tensor &hidden_states,
+    infinicore::Tensor &residual) {
+    input_layernorm_->forward_inplace(hidden_states, residual);
+    if ("linear_attention" == layer_type_) {
+        hidden_states = linear_attn_->forward(hidden_states);
+    } else {
+        hidden_states = self_attn_->forward(positions, hidden_states);
+    }
+
+    post_attention_layernorm_->forward_inplace(hidden_states, residual);
+    hidden_states = mlp_->forward(hidden_states);
+    return std::make_tuple(hidden_states, residual);
+}
+
+infinicore::Tensor Qwen35MoeDecoderLayer::forward(
+    const infinicore::Tensor &positions,
+    infinicore::Tensor &hidden_states) {
+    auto residual = hidden_states;
+    hidden_states = input_layernorm_->forward(hidden_states);
+    if ("linear_attention" == layer_type_) {
+        hidden_states = linear_attn_->forward(hidden_states);
+    } else {
+        hidden_states = self_attn_->forward(positions, hidden_states);
+    }
+    hidden_states = infinicore::op::add(residual, hidden_states);
+
+    residual = hidden_states;
+    hidden_states = post_attention_layernorm_->forward(hidden_states);
+    hidden_states = mlp_->forward(hidden_states);
+    return infinicore::op::add(residual, hidden_states);
+}
+
+} // namespace infinilm::models::qwen3_5_moe