feat: support Ktransformers, CPU-GPU MoE offload via FusedMoE layer

csrc/layers/moe/experts/fused_moe_experts.cpp

@@ -12,6 +12,11 @@ FusedMoeExperts::FusedMoeExperts(std::shared_ptr<infinilm::config::ModelConfig>
     num_experts_ = model_config->get<size_t>("num_experts");
     hidden_size_ = model_config->get<size_t>("hidden_size");
     const size_t intermediate_size = model_config->get<size_t>("moe_intermediate_size");
+    if (model_config->get_or<bool>("use_kt_moe", false)) {
+        // KT offload: expert weights are filtered out at load time; skip
+        // allocating GPU-side packed weights entirely.
+        return;
+    }
     const auto dtype = model_config->get_dtype();
     ASSERT(num_experts_ > 0);
 

csrc/layers/moe/fused_moe.cpp

@@ -1,5 +1,7 @@
 #include "fused_moe.hpp"
 
+#include "../moe/kt_moe_callback.hpp"
+
 #include "dispatcher/dispatcher_factory.hpp"
 #include "ep/ep_config.hpp"
 #include "runner/cuda_fused_moe_runner.hpp"
@@ -12,8 +14,15 @@ namespace infinilm::layers::moe {
 
 FusedMoE::FusedMoE(std::shared_ptr<infinilm::config::ModelConfig> model_config,
                    const infinicore::Device &device,
-                   size_t layer_id) {
-    (void)layer_id;
+                   size_t layer_id)
+    : layer_id_(layer_id),
+      skip_experts_(model_config->get_or<bool>("use_kt_moe", false)) {
+    if (skip_experts_) {
+        // KT offload: routed experts live on CPU (kt-kernel); skip building
+        // the GPU dispatcher/runner entirely. forward() consults the KT
+        // callback registry.
+        return;
+    }
 
     const EPConfig ep_config = make_ep_config();
     const size_t num_experts = model_config->get<size_t>("num_experts");
@@ -41,6 +50,22 @@ FusedMoE::FusedMoE(std::shared_ptr<infinilm::config::ModelConfig> model_config,
 infinicore::Tensor FusedMoE::forward(const infinicore::Tensor &hidden_states,
                                      const TopKOutput &topk_output,
                                      const MoeWeights &weights) const {
+    // KT (KTransformers) branch: delegate routed-expert compute to CPU-GPU
+    // heterogeneous offload before touching GPU weights.
+    {
+        auto kt_cb = infinilm::layers::moe::KTMoECallbackRegistry::instance().get(
+            static_cast<int>(layer_id_));
+        if (kt_cb) {
+            return (*kt_cb)(hidden_states, topk_output.topk_weights, topk_output.topk_ids,
+                            static_cast<int>(layer_id_));
+        }
+        if (skip_experts_) {
+            throw std::runtime_error(
+                "FusedMoE: use_kt_moe is enabled but no KT callback is registered for layer "
+                + std::to_string(layer_id_));
+        }
+    }
+
     auto dispatch_output = dispatcher_->dispatch(hidden_states, topk_output, workspace_);
     auto combine_input = runner_->run(dispatch_output, weights, workspace_);
     return dispatcher_->combine(combine_input, workspace_);

csrc/layers/moe/fused_moe.hpp

@@ -25,6 +25,8 @@ class FusedMoE final : public infinicore::nn::Module {
     std::shared_ptr<BaseDispatcher> dispatcher_;
     std::shared_ptr<MoeRunnerCore> runner_;
     mutable MoeWorkspace workspace_;
+    size_t layer_id_{0};
+    bool skip_experts_{false};
 };
 
 } // namespace infinilm::layers::moe

csrc/layers/moe/kt_moe_callback.hpp

@@ -0,0 +1,59 @@
+#pragma once
+#include "infinicore/tensor.hpp"
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <unordered_map>
+
+namespace infinilm::layers::moe {
+
+// Global registry of KTransformers MoE callbacks (one per layer_idx).
+//
+// Concurrency contract:
+//  - Callbacks are stored as immutable shared_ptr entries. get() copies the
+//    entry out under the lock and the user callback is invoked WITHOUT any
+//    lock held, so a callback that acquires the GIL can never deadlock
+//    against set()/clear() called from a GIL-holding thread.
+//  - get() is a single atomic lookup: no has()+call() TOCTOU window.
+class KTMoECallbackRegistry {
+public:
+    using CallbackFn = std::function<infinicore::Tensor(
+        const infinicore::Tensor &hidden_states,
+        const infinicore::Tensor &topk_weights,
+        const infinicore::Tensor &topk_ids,
+        int layer_idx)>;
+
+    static KTMoECallbackRegistry &instance() {
+        static KTMoECallbackRegistry reg;
+        return reg;
+    }
+
+    void set(int layer_idx, CallbackFn cb) {
+        std::lock_guard<std::mutex> lk(mtx_);
+        cbs_[layer_idx] = std::make_shared<const CallbackFn>(std::move(cb));
+    }
+
+    void clear() {
+        std::lock_guard<std::mutex> lk(mtx_);
+        cbs_.clear();
+    }
+
+    // Returns nullptr when no callback is registered for this layer.
+    // The returned pointer stays valid even if set/clear run concurrently.
+    std::shared_ptr<const CallbackFn> get(int layer_idx) {
+        std::lock_guard<std::mutex> lk(mtx_);
+        auto it = cbs_.find(layer_idx);
+        return it == cbs_.end() ? nullptr : it->second;
+    }
+
+    bool empty() {
+        std::lock_guard<std::mutex> lk(mtx_);
+        return cbs_.empty();
+    }
+
+private:
+    std::mutex mtx_;
+    std::unordered_map<int, std::shared_ptr<const CallbackFn>> cbs_;
+};
+
+} // namespace infinilm::layers::moe

csrc/models/deepseek/deepseek_decoder_layer.cpp

@@ -22,7 +22,7 @@ DeepseekDecoderLayer::DeepseekDecoderLayer(std::shared_ptr<infinilm::config::Mod
 
     if (use_moe) {
         mlp_ = std::make_shared<DeepseekMLP>(
-            this->register_module<deepseek_v2::DeepseekV2MoE>("mlp", model_config, device));
+            this->register_module<deepseek_v2::DeepseekV2MoE>("mlp", model_config, layer_idx, device));
     } else {
         mlp_ = std::make_shared<DeepseekMLP>(
             this->register_module<deepseek_v2::DeepseekV2MLP>("mlp", model_config, device));

csrc/models/deepseek_v2/deepseek_v2_decoder_layer.cpp

@@ -24,7 +24,7 @@ DeepseekV2DecoderLayer::DeepseekV2DecoderLayer(std::shared_ptr<infinilm::config:
             && layer_idx >= first_k_dense_replace
             && (moe_layer_freq == 0 || layer_idx % moe_layer_freq == 0);
     if (use_moe_) {
-        moe_mlp_ = this->register_module<DeepseekV2MoE>("mlp", model_config, device);
+        moe_mlp_ = this->register_module<DeepseekV2MoE>("mlp", model_config, layer_idx, device);
     } else {
         dense_mlp_ = this->register_module<DeepseekV2MLP>("mlp", model_config, device);
     }

csrc/models/deepseek_v2/deepseek_v2_moe.cpp

@@ -1,4 +1,5 @@
 #include "deepseek_v2_moe.hpp"
+#include "../../layers/moe/kt_moe_callback.hpp"
 
 #include "../../global_state/global_state.hpp"
 #include "../../utils.hpp"
@@ -131,9 +132,14 @@ infinicore::Tensor DeepseekV2Experts::forward(const infinicore::Tensor &hidden_s
 }
 
 DeepseekV2MoE::DeepseekV2MoE(std::shared_ptr<infinilm::config::ModelConfig> model_config,
-                             const infinicore::Device &device) {
+                             size_t layer_idx,
+                             const infinicore::Device &device)
+    : layer_idx_(layer_idx) {
+    skip_experts_ = model_config->get_or<bool>("use_kt_moe", false);
     INFINICORE_NN_MODULE_INIT(gate, model_config, device);
-    INFINICORE_NN_MODULE_INIT(experts, model_config, device);
+    if (!skip_experts_) {
+        INFINICORE_NN_MODULE_INIT(experts, model_config, device);
+    }
 
     const size_t n_shared_experts = model_config->get_or<size_t>("n_shared_experts", 0);
     has_shared_experts_ = n_shared_experts > 0;
@@ -151,7 +157,26 @@ infinicore::Tensor DeepseekV2MoE::forward(const infinicore::Tensor &hidden_state
     auto hidden_states_reshaped = hidden_states->view({shape[0] * shape[1], shape[2]});
 
     auto [routing_weights, selected_experts] = gate_->forward(hidden_states_reshaped);
-    auto final_hidden_states = experts_->forward(hidden_states_reshaped, selected_experts, routing_weights)->view(shape);
+
+    const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info();
+    infinicore::Tensor expert_output;
+    auto kt_cb = infinilm::layers::moe::KTMoECallbackRegistry::instance().get(static_cast<int>(layer_idx_));
+    if (kt_cb) {
+        if (rank_info.tp_size > 1) {
+            // Each rank would run KT on the full expert weights and the native path's
+            // partial-sum allreduce semantics do not apply -> explicit refusal.
+            throw std::runtime_error(
+                "DeepseekV2MoE: KT offload does not support tensor_parallel_size > 1");
+        }
+        expert_output = (*kt_cb)(hidden_states_reshaped, routing_weights, selected_experts, static_cast<int>(layer_idx_));
+    } else if (skip_experts_) {
+        throw std::runtime_error(
+            "DeepseekV2MoE: use_kt_moe is enabled but no KT callback is registered for layer "
+            + std::to_string(layer_idx_));
+    } else {
+        expert_output = experts_->forward(hidden_states_reshaped, selected_experts, routing_weights);
+    }
+    auto final_hidden_states = expert_output->view(shape);
     if (has_shared_experts_) {
         auto shared_out = shared_experts_->forward(hidden_states);
         final_hidden_states = infinicore::op::add(final_hidden_states, shared_out);

csrc/models/deepseek_v2/deepseek_v2_moe.hpp

@@ -63,6 +63,7 @@ class DeepseekV2Experts : public infinicore::nn::Module {
 class DeepseekV2MoE : public infinicore::nn::Module {
 public:
     DeepseekV2MoE(std::shared_ptr<infinilm::config::ModelConfig> model_config,
+                  size_t layer_idx,
                   const infinicore::Device &device);
 
     infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const;
@@ -72,6 +73,8 @@ class DeepseekV2MoE : public infinicore::nn::Module {
     INFINICORE_NN_MODULE(DeepseekV2Experts, experts);
     INFINICORE_NN_MODULE(DeepseekV2MLP, shared_experts);
     bool has_shared_experts_{false};
+    size_t layer_idx_{0};
+    bool skip_experts_{false};
 };
 
 } // namespace infinilm::models::deepseek_v2

csrc/models/qwen3_moe/qwen3_moe_sparse_moe_block.cpp

@@ -43,6 +43,9 @@ infinicore::Tensor Qwen3MoeSparseMoeBlock::forward(const infinicore::Tensor &hid
         infinicore::Tensor(),
     };
 
+    // KT (KTransformers) offload is handled inside FusedMoE::forward when
+    // a KT callback is registered for this layer.
+
     auto final_hidden_states = fused_moe_->forward(
         hidden_states_reshaped,
         topk_output,

csrc/models/qwen3_next/qwen3_next_sparse_moe_block.cpp

@@ -64,6 +64,7 @@ Qwen3NextSparseMoeBlock::Qwen3NextSparseMoeBlock(std::shared_ptr<infinilm::confi
     gate_ = this->register_module<infinilm::layers::moe::TopKRouter>("gate", model_config, device);
     experts_ = this->register_module<infinilm::layers::moe::FusedMoeExperts>("experts", model_config, device);
     fused_moe_ = this->register_module<infinilm::layers::moe::FusedMoE>("fused_moe", model_config, device, layer_idx);
+    (void)layer_idx;
     shared_expert_ = this->register_module<Qwen3NextSharedExpert>("shared_expert", model_config, device);
     shared_expert_gate_ = this->register_module<infinilm::layers::linear::ReplicatedLinear>(
         "shared_expert_gate",
@@ -86,6 +87,8 @@ infinicore::Tensor Qwen3NextSparseMoeBlock::forward(const infinicore::Tensor &hi
         selected_experts,
         infinicore::Tensor(),
     };
+    // KT (KTransformers) offload is handled inside FusedMoE::forward when
+    // a KT callback is registered for this layer.
     auto routed_states = fused_moe_->forward(
         hidden_states_reshaped,
         topk_output,

csrc/pybind11/bindings.cc

@@ -1,7 +1,10 @@
+#include "../layers/moe/kt_moe_callback.hpp"
 #include <pybind11/pybind11.h>
 
 #include "cache/cache.hpp"
 #include "engine/engine.hpp"
+#include <stdexcept>
+#include <string>
 
 namespace py = pybind11;
 
@@ -12,4 +15,40 @@ PYBIND11_MODULE(_infinilm, m) {
     infinilm::engine::bind_hook_registry(m);
     infinilm::engine::distributed::bind_dist_config(m);
     infinilm::engine::bind_infer_engine(m);
+
+    // ---- KTransformers MoE offload integration ----
+    // Register a Python callback (per layer) that receives
+    // (hidden, routing_weights, topk_ids, layer_idx) as infinicore tensors
+    // and must return the routed-expert output tensor (infinicore view).
+    m.def(
+        "set_kt_moe_callback", [](int layer_idx, py::function callback) {
+            infinilm::layers::moe::KTMoECallbackRegistry::instance().set(layer_idx,
+                                                                         [callback](const infinicore::Tensor &h, const infinicore::Tensor &w,
+                                                                                    const infinicore::Tensor &i, int l) -> infinicore::Tensor {
+                                                                             // Translate Python exceptions at the boundary while the GIL is
+                                                                             // held: the C++ worker thread has no GIL and no Python frame to
+                                                                             // unwind into, and py::error_already_set must not escape it.
+                                                                             py::gil_scoped_acquire gil;
+                                                                             try {
+                                                                                 return callback(h, w, i, l).cast<infinicore::Tensor>();
+                                                                             } catch (const py::error_already_set &e) {
+                                                                                 throw std::runtime_error(
+                                                                                     "KT MoE callback (layer " + std::to_string(l) + ") failed: " + e.what());
+                                                                             }
+                                                                         });
+        },
+        py::arg("layer_idx"), py::arg("callback"), "Register a KTransformers MoE callback for a given layer.");
+
+    m.def(
+        "clear_kt_moe_callbacks", []() {
+            infinilm::layers::moe::KTMoECallbackRegistry::instance().clear();
+        },
+        "Clear all KT MoE callbacks.");
+
+    // Release registered py::functions at module teardown (while the
+    // interpreter is still alive) instead of process-exit static destruction.
+    m.add_object("_kt_moe_cleanup",
+                 py::capsule(reinterpret_cast<void *>(1), "_kt_moe_cleanup", [](void *) {
+                     infinilm::layers::moe::KTMoECallbackRegistry::instance().clear();
+                 }));
 }