feat: support MOSS split QKV weights and batch host loading (#2160)

src/model_manager.cpp

@@ -1022,28 +1022,60 @@ bool ModelManager::load_tensor_parts(TensorState& state) {
 
 bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
     using ReadGroup = std::pair<ModelLoader::FileId, SDVersion>;
-    using ReadBatch = std::map<std::string, std::vector<TensorState*>>;
+    using ReadBatch = std::map<std::string, std::vector<ggml_tensor*>>;
     std::map<ReadGroup, std::vector<ReadBatch>> groups;
-    for (auto* state : states) {
-        if (state == nullptr)
-            continue;
-        if (state->sources.size() > 1) {
-            if (!load_tensor_parts(*state)) {
-                return false;
-            }
-            continue;
+    size_t part_count = 0;
+    for (const auto* state : states) {
+        if (state != nullptr && state->sources.size() > 1 && ggml_backend_buffer_is_host(state->tensor->buffer)) {
+            part_count += state->sources.size();
         }
-        auto& batches = groups[{state->source_file, state->source_version}];
-        // The loader supplies one destination per name; only conflicting types need another batch.
+    }
+    auto parts_ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
+        part_count > 0 ? ggml_init({part_count * ggml_tensor_overhead(), nullptr, true}) : nullptr, ggml_free);
+    if (part_count > 0 && !parts_ctx) {
+        return false;
+    }
+    auto add_destination = [&](const TensorState& state, const std::string& name, ggml_tensor* tensor) {
+        auto& batches = groups[{state.source_file, state.source_version}];
+        // The loader supplies one destination per name; incompatible destinations need another batch.
         auto batch = std::find_if(batches.begin(), batches.end(), [&](const ReadBatch& candidate) {
-            auto found = candidate.find(state->name);
-            return found == candidate.end() || found->second.front()->tensor->type == state->tensor->type;
+            auto found = candidate.find(name);
+            return found == candidate.end() ||
+                   (found->second.front()->type == tensor->type && ggml_are_same_shape(found->second.front(), tensor));
         });
         if (batch == batches.end()) {
             batches.emplace_back();
             batch = std::prev(batches.end());
         }
-        (*batch)[state->name].push_back(state);
+        (*batch)[name].push_back(tensor);
+    };
+    for (auto* state : states) {
+        if (state == nullptr)
+            continue;
+        if (state->sources.size() > 1) {
+            if (ggml_backend_buffer_is_host(state->tensor->buffer)) {
+                const size_t size = ggml_nbytes(state->tensor);
+                size_t offset     = 0;
+                for (const auto& source : state->sources) {
+                    auto* part             = ggml_new_tensor(parts_ctx.get(), state->tensor->type, source.n_dims, source.ne);
+                    const size_t part_size = ggml_nbytes(part);
+                    if (part_size > size - offset) {
+                        return false;
+                    }
+                    part->data   = static_cast<uint8_t*>(state->tensor->data) + offset;
+                    part->buffer = state->tensor->buffer;
+                    add_destination(*state, source.name, part);
+                    offset += part_size;
+                }
+                if (offset != size) {
+                    return false;
+                }
+            } else if (!load_tensor_parts(*state)) {
+                return false;
+            }
+            continue;
+        }
+        add_destination(*state, state->name, state->tensor);
     }
     for (auto& group : groups) {
         for (auto& batch : group.second) {
@@ -1057,7 +1089,7 @@ bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
                 auto found = batch.find(source.name);
                 if (found == batch.end())
                     return true;
-                *dst = found->second.front()->tensor;
+                *dst = found->second.front();
                 std::lock_guard<std::mutex> lock(mutex);
                 loaded.insert(source.name);
                 return true;
@@ -1068,15 +1100,18 @@ bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
             if (!success || loaded != names)
                 return false;
             for (auto& entry : batch) {
-                auto* first = entry.second.front()->tensor;
-                for (auto* state : entry.second) {
-                    if (state->tensor != first)
-                        ggml_backend_tensor_copy(first, state->tensor);
-                    state->loaded_to_params_backend = true;
+                auto* first = entry.second.front();
+                for (auto* tensor : entry.second) {
+                    if (tensor != first)
+                        ggml_backend_tensor_copy(first, tensor);
                 }
             }
         }
     }
+    for (auto* state : states) {
+        if (state != nullptr)
+            state->loaded_to_params_backend = true;
+    }
     return true;
 }
 

src/name_conversion.cpp

@@ -1013,6 +1013,16 @@ static std::string convert_diffusers_dit_to_original_pixart(std::string name) {
     return prefix + name;
 }
 
+static std::string convert_moss_image_dit_name(std::string name) {
+    static const std::vector<std::pair<std::string, std::string>> name_map = {
+        {".proj.q.weight", ".proj.qkv.weight"},
+        {".proj.k.weight", ".proj.qkv.weight.1"},
+        {".proj.v.weight", ".proj.qkv.weight.2"},
+    };
+    replace_with_name_map(name, name_map);
+    return name;
+}
+
 static std::string convert_ming_image_dit_name(std::string name) {
     static const std::vector<std::pair<std::string, std::string>> name_map = {
         {"all_x_embedder.2-1.", "x_embedder."},
@@ -1053,6 +1063,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
         name = convert_diffusers_dit_to_original_flux(name);
     } else if (sd_version_is_hunyuan_video(version)) {
         name = convert_hunyuan_video_to_original_flux(name);
+    } else if (version == VERSION_MOSS_IMAGE) {
+        name = convert_moss_image_dit_name(name);
     } else if (version == VERSION_MING_IMAGE) {
         name = convert_ming_image_dit_name(name);
     } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version)) {