feat: support for cross-device row split (#1735)

src/core/ggml_extend_backend.cpp

@@ -665,8 +665,9 @@ SDBackendManager::~SDBackendManager() {
 
 void SDBackendManager::reset() {
     backends_.clear();
-    runtime_assignment_ = {};
-    params_assignment_  = {};
+    runtime_assignment_    = {};
+    params_assignment_     = {};
+    split_mode_assignment_ = {};
 }
 
 static std::vector<std::string> split_device_list(const std::string& value) {
@@ -758,6 +759,7 @@ bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule modu
 
 bool SDBackendManager::init(const char* backend_spec,
                             const char* params_backend_spec,
+                            const char* split_mode_spec,
                             std::string* error) {
     reset();
 
@@ -767,10 +769,51 @@ bool SDBackendManager::init(const char* backend_spec,
     if (!sd_parse_backend_assignment(SAFE_STR(params_backend_spec), &params_assignment_, error)) {
         return false;
     }
+    if (!sd_parse_backend_assignment(SAFE_STR(split_mode_spec), &split_mode_assignment_, error)) {
+        return false;
+    }
 
     return validate(error);
 }
 
+SDSplitMode SDBackendManager::split_mode(SDBackendModule module) const {
+    return lower_copy(trim_copy(split_mode_assignment_.get(module))) == "row" ? SDSplitMode::ROW
+                                                                              : SDSplitMode::LAYER;
+}
+
+ggml_backend_buffer_type_t SDBackendManager::split_buffer_type(ggml_backend_t backend,
+                                                               const std::vector<float>& tensor_split) {
+    if (backend == nullptr) {
+        return nullptr;
+    }
+    ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+    if (dev == nullptr) {
+        return nullptr;
+    }
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+    if (reg == nullptr) {
+        return nullptr;
+    }
+    auto fn = (ggml_backend_split_buffer_type_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_split_buffer_type");
+    if (fn == nullptr) {
+        return nullptr;
+    }
+    int main_device        = -1;
+    const size_t dev_count = ggml_backend_reg_dev_count(reg);
+    for (size_t i = 0; i < dev_count; ++i) {
+        if (ggml_backend_reg_dev_get(reg, i) == dev) {
+            main_device = (int)i;
+            break;
+        }
+    }
+    if (main_device < 0) {
+        return nullptr;
+    }
+    std::vector<float> padded_split(std::max<size_t>(tensor_split.size(), 64), 0.0f);
+    std::copy(tensor_split.begin(), tensor_split.end(), padded_split.begin());
+    return fn(main_device, padded_split.data());
+}
+
 bool SDBackendManager::validate(std::string* error) const {
     auto validate_single_runtime_name = [&](const std::string& name) -> bool {
         if (is_default_backend_token(name)) {
@@ -826,9 +869,20 @@ bool SDBackendManager::validate(std::string* error) const {
         }
         return validate_single_runtime_name(name);
     };
+    auto validate_split_mode_name = [&](const std::string& name) -> bool {
+        const std::string lower = lower_copy(trim_copy(name));
+        if (lower.empty() || lower == "layer" || lower == "row") {
+            return true;
+        }
+        if (error != nullptr) {
+            *error = "invalid split mode '" + name + "' (expected layer or row)";
+        }
+        return false;
+    };
 
     if (!validate_runtime_name(runtime_assignment_.default_name) ||
-        !validate_params_name(params_assignment_.default_name)) {
+        !validate_params_name(params_assignment_.default_name) ||
+        !validate_split_mode_name(split_mode_assignment_.default_name)) {
         return false;
     }
     for (const auto& kv : runtime_assignment_.module_names) {
@@ -841,6 +895,11 @@ bool SDBackendManager::validate(std::string* error) const {
             return false;
         }
     }
+    for (const auto& kv : split_mode_assignment_.module_names) {
+        if (!validate_split_mode_name(kv.second)) {
+            return false;
+        }
+    }
     return true;
 }
 

src/core/ggml_extend_backend.h

@@ -38,10 +38,16 @@ struct SDBackendHandleDeleter {
 
 using SDBackendHandle = std::unique_ptr<struct ggml_backend, SDBackendHandleDeleter>;
 
+enum class SDSplitMode {
+    LAYER,
+    ROW,
+};
+
 class SDBackendManager {
 private:
     SDBackendAssignment runtime_assignment_;
     SDBackendAssignment params_assignment_;
+    SDBackendAssignment split_mode_assignment_;
     std::unordered_map<std::string, SDBackendHandle> backends_;
 
 public:
@@ -53,6 +59,7 @@ public:
 
     bool init(const char* backend_spec,
               const char* params_backend_spec,
+              const char* split_mode_spec,
               std::string* error);
     void reset();
 
@@ -61,6 +68,10 @@ public:
 
     std::vector<ggml_backend_t> runtime_backends(SDBackendModule module);
 
+    SDSplitMode split_mode(SDBackendModule module) const;
+    ggml_backend_buffer_type_t split_buffer_type(ggml_backend_t backend,
+                                                 const std::vector<float>& tensor_split);
+
     bool runtime_backend_is_cpu(SDBackendModule module);
     bool params_backend_is_cpu(SDBackendModule module);
     bool params_backend_is_disk(SDBackendModule module) const;

src/core/layer_split_partition.cpp

@@ -9,7 +9,7 @@
 
 namespace sd {
 
-    static int tensor_block_index(const std::string& name) {
+    int layer_split_tensor_block_index(const std::string& name) {
         static const char* block_keywords[] = {"transformer_blocks.", "joint_blocks.", "double_blocks.",
                                                "single_blocks.", "blocks.", "block.", "layers."};
         for (const char* keyword : block_keywords) {
@@ -92,7 +92,7 @@ namespace sd {
         int n_blocks              = 0;
         for (const auto& kv : tensors) {
             int64_t bytes = (int64_t)ggml_nbytes(kv.second);
-            int idx       = split_tensors.count(kv.first) != 0 ? tensor_block_index(kv.first) : -1;
+            int idx       = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
             if (idx >= 0) {
                 block_bytes[idx] += bytes;
                 total_block_bytes += bytes;
@@ -160,7 +160,7 @@ namespace sd {
 
         for (const auto& kv : tensors) {
             size_t target = 0;
-            int idx       = split_tensors.count(kv.first) != 0 ? tensor_block_index(kv.first) : -1;
+            int idx       = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
             if (idx >= 0) {
                 while (target < boundaries.size() && idx >= boundaries[target]) {
                     target++;

src/core/layer_split_partition.h

@@ -11,6 +11,7 @@
 namespace sd {
 
     std::string layer_split_backend_device_display_name(ggml_backend_t backend);
+    int layer_split_tensor_block_index(const std::string& name);
 
     std::vector<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
         const std::string& desc,