[2026春季][T2-1-3]xuzheng567

csrc/models/qwen3_moe/qwen3_moe_experts.cpp

@@ -1,64 +1,65 @@
 #include "qwen3_moe_experts.hpp"
 
+#include "../../config/model_config.hpp"
+#include "../../global_state/global_state.hpp"
 #include "infinicore/ops.hpp"
 
+#include <optional>
 #include <string>
 
 namespace infinilm::models::qwen3_moe {
 
 Qwen3MoeExperts::Qwen3MoeExperts(std::shared_ptr<infinilm::config::ModelConfig> model_config,
                                  const infinicore::Device &device) {
+    const auto &dtype = model_config->get_dtype();
+    const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info();
+    const int tp_rank = rank_info.tp_rank;
+    const int tp_size = rank_info.tp_size;
 
     num_experts_ = model_config->get<size_t>("num_experts");
     num_experts_per_tok_ = model_config->get<size_t>("num_experts_per_tok");
+    hidden_size_ = model_config->get<size_t>("hidden_size");
+    const size_t intermediate_size = model_config->get<size_t>("moe_intermediate_size");
 
     ASSERT((num_experts_ > 0) && (num_experts_per_tok_ > 0) && (num_experts_per_tok_ <= num_experts_));
-
-    for (size_t i = 0; i < num_experts_; ++i) {
-        experts_.push_back(this->register_module<Qwen3MoeMLP>(std::to_string(i), model_config, device));
+    ASSERT(intermediate_size % static_cast<size_t>(tp_size) == 0);
+    intermediate_size_per_partition_ = intermediate_size / static_cast<size_t>(tp_size);
+
+    INFINICORE_NN_PARAMETER_INIT(w1, ({num_experts_, 2 * intermediate_size_per_partition_, hidden_size_}, dtype, device));
+    INFINICORE_NN_PARAMETER_INIT(w2, ({num_experts_, hidden_size_, intermediate_size_per_partition_}, dtype, device));
+
+    for (size_t expert = 0; expert < num_experts_; ++expert) {
+        auto gate_weight = w1_
+                               ->narrow({{0, expert, 1}, {1, 0, intermediate_size_per_partition_}})
+                               ->squeeze(0);
+        auto up_weight = w1_
+                             ->narrow({{0, expert, 1}, {1, intermediate_size_per_partition_, intermediate_size_per_partition_}})
+                             ->squeeze(0);
+        auto down_weight = w2_
+                               ->narrow({{0, expert, 1}})
+                               ->squeeze(0);
+
+        const std::string prefix = std::to_string(expert) + ".";
+        this->register_parameter(
+            prefix + "gate_proj.weight",
+            infinicore::nn::Parameter(gate_weight, 0, tp_rank, tp_size));
+        this->register_parameter(
+            prefix + "up_proj.weight",
+            infinicore::nn::Parameter(up_weight, 0, tp_rank, tp_size));
+        this->register_parameter(
+            prefix + "down_proj.weight",
+            infinicore::nn::Parameter(down_weight, 1, tp_rank, tp_size));
     }
 }
 
 infinicore::Tensor Qwen3MoeExperts::forward(const infinicore::Tensor &hidden_states,
                                             const infinicore::Tensor &top_k_index,
                                             const infinicore::Tensor &top_k_weights) const {
     ASSERT(hidden_states->ndim() == 2);
+    ASSERT(top_k_index->ndim() == 2 && top_k_weights->ndim() == 2);
 
-    auto top_k_weights_cpu = top_k_weights->to(infinicore::Device::Type::CPU);
-    auto top_k_index_cpu = top_k_index->to(infinicore::Device::Type::CPU);
-
-    int *top_k_index_ptr = reinterpret_cast<int *>(top_k_index_cpu->data());
-    float *top_k_weights_ptr = reinterpret_cast<float *>(top_k_weights_cpu->data());
-
-    size_t ntoken = hidden_states->shape()[0];
-    int index;
-    float score;
-
-    auto final_hidden_states = infinicore::Tensor::empty(hidden_states->shape(), hidden_states->dtype(), hidden_states->device());
-    for (size_t itok = 0; itok < ntoken; ++itok) {
-        auto hidden_states_i = hidden_states->narrow({{0, itok, 1}});
-        const size_t route_row = itok * num_experts_per_tok_;
-
-        infinicore::Tensor final_hidden_states_i;
-        for (size_t k = 0; k < num_experts_per_tok_; ++k) {
-            index = top_k_index_ptr[route_row + k];
-            score = top_k_weights_ptr[route_row + k];
-
-            ASSERT(index >= 0 && static_cast<size_t>(index) < num_experts_);
-
-            experts_[index]->set_alpha(score);
-            auto expert_out = experts_[index]->forward(hidden_states_i);
-
-            if (k == 0) {
-                final_hidden_states_i = expert_out;
-            } else {
-                infinicore::op::add_(final_hidden_states_i, final_hidden_states_i, expert_out);
-            }
-        }
-
-        final_hidden_states->narrow({{0, itok, 1}})->copy_from(final_hidden_states_i);
-    }
-    return final_hidden_states;
+    return infinicore::op::fused_moe(hidden_states, top_k_index, top_k_weights, w1_, w2_, std::nullopt, std::nullopt,
+                                     infinicore::op::FusedMoeActivation::Swiglu);
 }
 
 } // namespace infinilm::models::qwen3_moe

csrc/models/qwen3_moe/qwen3_moe_experts.hpp

@@ -1,15 +1,17 @@
 #pragma once
 
-#include "../../layers/moe/legacy/moe_mlp.hpp"
 #include "infinicore/nn/module.hpp"
+#include "infinicore/nn/parameter.hpp"
 #include "infinicore/tensor.hpp"
 
 #include <cstddef>
 #include <memory>
 
-namespace infinilm::models::qwen3_moe {
+namespace infinilm::config {
+class ModelConfig;
+}
 
-using Qwen3MoeMLP = infinilm::layers::moe::legacy::MoeMLP;
+namespace infinilm::models::qwen3_moe {
 
 class Qwen3MoeExperts : public infinicore::nn::Module {
 public:
@@ -21,9 +23,13 @@ class Qwen3MoeExperts : public infinicore::nn::Module {
                                const infinicore::Tensor &top_k_weights) const;
 
 protected:
-    INFINICORE_NN_MODULE_VEC(Qwen3MoeMLP, experts);
+    INFINICORE_NN_PARAMETER(w1);
+    INFINICORE_NN_PARAMETER(w2);
+
     size_t num_experts_per_tok_{0};
     size_t num_experts_{0};
+    size_t hidden_size_{0};
+    size_t intermediate_size_per_partition_{0};
 };
 
 } // namespace infinilm::models::qwen3_moe

python/infinilm/modeling_utils.py

@@ -59,8 +59,11 @@ def _is_internal_moe_packed_weight(key: str) -> bool:
     # InfiniLM registers packed MoE parameters internally. HF checkpoints
     # provide per-expert gate/up/down weights instead, so these packed tensors
     # are expected missing keys during non-strict checkpoint loading.
-    return key.endswith(".mlp.experts.w13_weight") or key.endswith(
-        ".mlp.experts.w2_weight"
+    return (
+        key.endswith(".mlp.experts.w13_weight")
+        or key.endswith(".mlp.experts.w2_weight")
+        or key.endswith(".mlp.experts.w1")
+        or key.endswith(".mlp.experts.w2")
     )