feat: add MoE inference and expert parallel support

Add reusable MoE routing, dispatch, expert, and runner layers for sparse MoE models.

Wire Qwen3/DeepSeek MoE implementations to the shared MoE path, expose EP backend configuration through DistConfig-backed bench/server/LLM entrypoints, and keep graph replay output handling compatible with MoE execution.

Author: qinyiqun

csrc/config/model_config.hpp

@@ -5,8 +5,10 @@
 #include "infinicore/ops.hpp"
 #include "quant_config.hpp"
 #include <fstream>
+#include <initializer_list>
 #include <iostream>
 #include <string>
+#include <string_view>
 #include <vector>
 
 namespace infinilm::config {
@@ -38,16 +40,33 @@ class ModelConfig {
 
     template <typename T>
     T get_or(const std::string &key, const T &default_value) const {
-        if (!config_json.contains(key) || config_json.at(key).is_null()) {
-            return default_value;
-        }
-        try {
-            return config_json.at(key).get<T>();
-        } catch (const nlohmann::json::type_error &) {
-            // If type conversion fails, return default value
-            return default_value;
+        return get_or<T>({std::string_view(key)}, default_value);
+    }
+
+    template <typename T>
+    T get_or(std::initializer_list<std::string_view> keys, const T &default_value) const {
+        for (std::string_view key : keys) {
+            if (key.empty()) {
+                continue;
+            }
+            auto it = config_json.find(std::string(key));
+            if (it == config_json.end() || it->is_null()) {
+                continue;
+            }
+            try {
+                return it->get<T>();
+            } catch (const nlohmann::json::type_error &) {
+                return default_value;
+            }
         }
+        return default_value;
     }
+
+    template <typename T>
+    T get_or_alias(const std::string &key, const std::string &alias, const T &default_value) const {
+        return get_or<T>({key, alias}, default_value);
+    }
+
     size_t get_kv_dim() const {
         return get<size_t>("hidden_size") * get<size_t>("num_key_value_heads") / get<size_t>("num_attention_heads");
     }

csrc/engine/compiler/paged_compiler.cpp

@@ -80,12 +80,15 @@ void PagedCompiler::compile() {
             auto input = make_decode_input(b);
 
             barrier_->wait();
+            (void)model_->forward(input);
+            infinicore::context::syncStream();
             // Capture must not start with stale Marlin locks from previous
             // warmup/capture attempts. This reset is intentionally outside
             // graph capture; the current implementation still pays a memset
             // before every graph replay in get_compiled().
             model_->reset_runtime_state();
             infinicore::context::syncStream();
+
             infinicore::context::startGraphRecording();
             auto output = model_->forward(input);
             auto graph = infinicore::context::stopGraphRecording();

csrc/engine/compiler/static_batching_compiler.cpp

@@ -29,14 +29,18 @@ void StaticBatchingCompiler::compile() {
         };
 
         barrier_->wait();
+        (void)model_->forward(input);
+        infinicore::context::syncStream();
+
         infinicore::context::startGraphRecording();
         auto output = model_->forward(input);
         auto graph = infinicore::context::stopGraphRecording();
         barrier_->wait();
 
         auto shared_output = std::shared_ptr<InfinilmModel::Output>(new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)});
+        auto replay_output = std::shared_ptr<InfinilmModel::Output>(new InfinilmModel::Output{shared_output->logits->resume_from_blob_()});
 
-        compiled_map_[std::make_tuple(b, 1)] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)};
+        compiled_map_[std::make_tuple(b, 1)] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output), replay_output};
     }
 }
 
@@ -56,7 +60,7 @@ StaticBatchingCompiler::Compiled StaticBatchingCompiler::get_compiled(
             graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value());
 
             auto graph = std::get<0>(result->second.compiled);
-            auto shared_output = std::shared_ptr<InfinilmModel::Output>(new InfinilmModel::Output{std::get<1>(result->second.compiled)->logits->resume_from_blob_()});
+            auto shared_output = result->second.replay_output;
             return std::make_tuple(graph, shared_output);
         }
     } else {

csrc/engine/compiler/static_batching_compiler.hpp

@@ -25,6 +25,9 @@ class StaticBatchingCompiler : public GraphCompiler {
     struct CompiledResult {
         InfinilmModel::Input input;
         Compiled compiled;
+        // Graph capture stores a GraphTensor in compiled. Replay returns a
+        // normal output handle restored from the same graph output blob.
+        std::shared_ptr<InfinilmModel::Output> replay_output;
     };
 
     std::unordered_map<

csrc/engine/distributed/dist_config.cpp

@@ -22,7 +22,7 @@ DistConfig::operator std::string() const {
             repr += ", ";
         }
     }
-    repr += "])";
+    repr += "], moe_ep_backend=" + moe_ep_backend + ", moe_ep_size=" + std::to_string(moe_ep_size) + ")";
     return repr;
 }
 

csrc/engine/distributed/dist_config.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <cstddef>
 #include <string>
 #include <vector>
 
@@ -8,6 +9,8 @@ namespace infinilm::engine::distributed {
 struct DistConfig {
     // Device IDs for each rank in tensor parallelism
     std::vector<int> tp_device_ids;
+    std::string moe_ep_backend{"disabled"};
+    size_t moe_ep_size{1};
 
     DistConfig();
     explicit DistConfig(int tp_size);

csrc/engine/infer_engine.cpp

@@ -3,6 +3,7 @@
 #include "spdlog/spdlog.h"
 #include <future>
 #include <stdexcept>
+#include <unordered_set>
 
 namespace infinilm::engine {
 
@@ -32,7 +33,12 @@ InferEngine::InferEngine(
 
     // Load model config if model_path is provided, model_path must be valid, and config.json exists
     this->model_config_ = infinilm::config::ConfigFactory::createConfig(config_str);
-    auto infinilm_config = std::make_shared<infinilm::global_state::InfinilmConfig>(attention_backend, this->model_config_, use_mla);
+    auto infinilm_config = std::make_shared<infinilm::global_state::InfinilmConfig>(
+        attention_backend,
+        this->model_config_,
+        use_mla,
+        distributed_config.moe_ep_backend,
+        distributed_config.moe_ep_size);
 
     // Only support offline int8 kv cache quantization in this version
     if (kv_cache_dtype.has_value()) {
@@ -67,19 +73,19 @@ void InferEngine::load_param(const std::string &name, const infinicore::Tensor &
     }
 }
 
-void InferEngine::load_params(const std::unordered_map<std::string, infinicore::Tensor> &params) {
+void InferEngine::load_params(const std::unordered_map<std::string, infinicore::Tensor> &params, bool strict) {
     if (workers_.size() <= 1 || weight_load_mode_ == "sync") {
         for (auto &worker : workers_) {
-            worker->load_params(params);
+            worker->load_params(params, strict);
         }
         return;
     }
 
     std::vector<std::future<void>> futures;
     futures.reserve(workers_.size());
     for (auto &worker : workers_) {
-        futures.emplace_back(std::async(std::launch::async, [&worker, &params] {
-            worker->load_params(params);
+        futures.emplace_back(std::async(std::launch::async, [&worker, &params, strict] {
+            worker->load_params(params, strict);
         }));
     }
     for (auto &future : futures) {
@@ -118,7 +124,17 @@ std::vector<std::string> InferEngine::state_dict_keys() {
     if (0 == workers_.size()) {
         throw std::runtime_error(" Model object not found. ");
     }
-    return workers_.front()->state_dict_keys();
+    std::vector<std::string> ordered_keys;
+    std::unordered_set<std::string> seen_keys;
+    for (auto &worker : workers_) {
+        for (const auto &key : worker->state_dict_keys()) {
+            // Preserve first-seen worker order while removing duplicate TP keys.
+            if (seen_keys.emplace(key).second) {
+                ordered_keys.push_back(key);
+            }
+        }
+    }
+    return ordered_keys;
 }
 
 //------------------------------------------------------

csrc/engine/infer_engine.hpp

@@ -8,6 +8,7 @@
 #include "rank_barrier.hpp"
 #include "rank_worker.hpp"
 
+#include <cstddef>
 #include <optional>
 #include <string>
 #include <unordered_map>
@@ -37,7 +38,7 @@ class InferEngine {
     void load_param(const std::string &name, const infinicore::Tensor &param);
 
     // Load a batch of parameters to all workers, syncing each worker once after the batch.
-    void load_params(const std::unordered_map<std::string, infinicore::Tensor> &params);
+    void load_params(const std::unordered_map<std::string, infinicore::Tensor> &params, bool strict = true);
 
     // process the weights after loading on all workers (e.g., for quantization)
     void process_weights_after_loading();

csrc/engine/rank_worker.cpp

@@ -87,14 +87,15 @@ void RankWorker::load_param(const std::string &name,
 //------------------------------------------------------
 // load_params -- synchronous batch load
 //------------------------------------------------------
-void RankWorker::load_params(const std::unordered_map<std::string, infinicore::Tensor> &params) {
+void RankWorker::load_params(const std::unordered_map<std::string, infinicore::Tensor> &params, bool strict) {
     {
         std::lock_guard<std::mutex> lock(mutex_);
         if (should_exit_) {
             throw std::runtime_error("RankWorker is closing; cannot load_params");
         }
 
         pending_params_ = params;
+        pending_params_strict_ = strict;
         job_cmd_ = Command::LOAD_BATCH;
         has_job_ = true;
         job_done_ = false;
@@ -295,6 +296,7 @@ void RankWorker::thread_loop() {
             std::string local_param_name;
             infinicore::Tensor local_param;
             std::unordered_map<std::string, infinicore::Tensor> local_params;
+            bool local_params_strict = true;
             Input local_args;
             std::unique_ptr<cache::CacheConfig> local_cache_config;
 
@@ -314,6 +316,10 @@ void RankWorker::thread_loop() {
                     local_param = pending_param_;
                 } else if (local_cmd == Command::LOAD_BATCH) {
                     local_params = std::move(pending_params_);
+                    // strict is copied with the batch because loading runs on
+                    // the worker thread after the caller releases the mutex.
+                    local_params_strict = pending_params_strict_;
+                    pending_params_strict_ = true;
                     pending_params_.clear();
                 } else if (local_cmd == Command::PREPROCESS) {
 
@@ -353,7 +359,7 @@ void RankWorker::thread_loop() {
 
             } else if (local_cmd == Command::LOAD_BATCH) {
                 try {
-                    model_->load_parameters_no_sync(local_params);
+                    model_->load_parameters_no_sync(local_params, local_params_strict);
                     infinicore::context::syncStream();
                 } catch (const std::exception &e) {
                     {

csrc/engine/rank_worker.hpp

@@ -87,7 +87,7 @@ class RankWorker {
     void load_param(const std::string &name,
                     const infinicore::Tensor &param);
 
-    void load_params(const std::unordered_map<std::string, infinicore::Tensor> &params);
+    void load_params(const std::unordered_map<std::string, infinicore::Tensor> &params, bool strict = true);
 
     void process_weights_after_loading();
 
@@ -150,6 +150,7 @@ class RankWorker {
     std::string pending_param_name_;
     infinicore::Tensor pending_param_;
     std::unordered_map<std::string, infinicore::Tensor> pending_params_;
+    bool pending_params_strict_ = true;
     Input pending_args_;
     std::unique_ptr<cache::CacheConfig> pending_cache_config_;