feat: support for cross-device row split

docs/backend.md

@@ -85,6 +85,36 @@ are disabled for it.
 
 Use `--list-devices` to see the device names available on the system.
 
+### Row split (`--split-mode row`)
+
+`--split-mode` selects how a multi-device module distributes its weights:
+`layer` (the default, described above) or `row`. It accepts a single mode or
+per-module assignments:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --backend "diffusion=cuda0&cuda1" --split-mode row
+sd-cli -m model.safetensors -p "a cat" --backend "diffusion=cuda0&cuda1,te=cuda0&cuda1" --split-mode diffusion=row,te=layer
+```
+
+In row mode the module keeps executing on its main (first listed) device, but
+its transformer-block matmul weights are allocated in the backend's row-split
+buffer type, which slices each weight's rows across the listed devices in
+proportion to free memory and runs those matmuls on all devices in parallel.
+Compared to a layer split this uses all GPUs within every layer (instead of
+sequentially device by device) at the cost of a cross-device reduction per
+matmul - usually the faster option when the devices have fast interconnect.
+
+Row split requires backend support for split buffers and is currently
+available on CUDA only; on other backends (or when the listed devices belong
+to different backend registries) the module falls back to a layer split.
+Embeddings, normalization weights, biases and other non-block tensors stay in
+regular buffers on the main device.
+
+Direct ("immediately") LoRA application cannot patch row-split tensors; with
+`--split-mode row` the automatic LoRA mode selects runtime application, and an
+explicit `--lora-apply-mode immediately` skips the split tensors with a
+warning.
+
 ## Modules
 
 | Module | Purpose | Accepted names |

examples/common/common.cpp

@@ -468,6 +468,13 @@ ArgOptions SDContextParams::get_options() {
          "parameter backend assignment, e.g. disk, cpu, or diffusion=disk,clip=cpu",
          (int)',',
          &params_backend},
+        {"",
+         "--split-mode",
+         "weight distribution for modules assigned multiple devices (--backend \"diffusion=cuda0&cuda1\"): "
+         "layer (whole transformer blocks per device, default) or row (matmul rows split across devices, CUDA only). "
+         "Accepts a single mode or per-module assignments, e.g. row or diffusion=row,te=layer",
+         (int)',',
+         &split_mode},
         {"",
          "--rpc-servers",
          "comma-separated list of RPC servers to connect to for offloading, in the format host:port, e.g. localhost:50052,192.168.1.3:50052",
@@ -830,6 +837,7 @@ std::string SDContextParams::to_string() const {
         << "  eager_load: " << (eager_load ? "true" : "false") << ",\n"
         << "  backend: \"" << backend << "\",\n"
         << "  params_backend: \"" << params_backend << "\",\n"
+        << "  split_mode: \"" << split_mode << "\",\n"
         << "  enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
         << "  control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n"
         << "  clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n"
@@ -910,6 +918,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.eager_load                      = eager_load;
     sd_ctx_params.backend                         = effective_backend.c_str();
     sd_ctx_params.params_backend                  = effective_params_backend.c_str();
+    sd_ctx_params.split_mode                      = split_mode.c_str();
     sd_ctx_params.rpc_servers                     = rpc_servers.c_str();
     return sd_ctx_params;
 }

examples/common/common.h

@@ -151,6 +151,7 @@ struct SDContextParams {
     bool eager_load             = false;
     std::string backend;
     std::string params_backend;
+    std::string split_mode;
     std::string rpc_servers;
     std::string effective_backend;
     std::string effective_params_backend;

include/stable-diffusion.h

@@ -227,6 +227,7 @@ typedef struct {
     bool eager_load;  // Load all params into the params backend at model-load time instead of lazily on first use
     const char* backend;
     const char* params_backend;
+    const char* split_mode;  // weight distribution for multi-device modules: layer (default) or row, or per-module assignments e.g. "diffusion=row"
     const char* rpc_servers;
 } sd_ctx_params_t;
 

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 @@ class SDBackendManager {
 
     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 @@ class SDBackendManager {
 
     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,

src/model_manager.cpp

@@ -100,12 +100,41 @@ size_t estimate_tensors_size(const std::map<std::string, ggml_tensor*>& tensors)
     return size;
 }
 
+void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft) {
+    if (compute_backend == nullptr) {
+        return;
+    }
+    if (split_buft == nullptr) {
+        split_buffer_types_.erase(compute_backend);
+        return;
+    }
+    split_buffer_types_[compute_backend] = split_buft;
+}
+
+bool ModelManager::tensor_shape_supports_split_buffer(const ggml_tensor* tensor) {
+    return tensor != nullptr &&
+           tensor->view_src == nullptr &&
+           ggml_is_contiguous(tensor) &&
+           ggml_n_dims(tensor) == 2 &&
+           tensor->ne[0] >= 256 &&
+           tensor->ne[1] >= 256;
+}
+
+ggml_backend_buffer_type_t ModelManager::split_buffer_type_for(const TensorState& state) const {
+    if (!state.allow_split_buffer || !tensor_shape_supports_split_buffer(state.tensor)) {
+        return nullptr;
+    }
+    auto it = split_buffer_types_.find(state.compute_backend);
+    return it != split_buffer_types_.end() ? it->second : nullptr;
+}
+
 bool ModelManager::register_param_tensors(const std::string& desc,
                                           std::map<std::string, ggml_tensor*> tensors,
                                           ResidencyMode residency_mode,
                                           ggml_backend_t compute_backend,
                                           ggml_backend_t params_backend,
-                                          size_t* registered_tensor_size) {
+                                          size_t* registered_tensor_size,
+                                          bool allow_split_buffer) {
     if (desc.empty()) {
         LOG_ERROR("model manager tensor desc is empty");
         return false;
@@ -129,13 +158,14 @@ bool ModelManager::register_param_tensors(const std::string& desc,
         }
         ggml_set_name(tensor, name.c_str());
 
-        auto state             = std::make_unique<TensorState>();
-        state->name            = name;
-        state->tensor          = tensor;
-        state->desc            = desc;
-        state->residency_mode  = residency_mode;
-        state->compute_backend = compute_backend;
-        state->params_backend  = params_backend;
+        auto state                = std::make_unique<TensorState>();
+        state->name               = name;
+        state->tensor             = tensor;
+        state->desc               = desc;
+        state->residency_mode     = residency_mode;
+        state->compute_backend    = compute_backend;
+        state->params_backend     = params_backend;
+        state->allow_split_buffer = allow_split_buffer;
         new_states.push_back(std::move(state));
     }
 
@@ -237,7 +267,7 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector<TensorState*
 }
 
 bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorState*>& states) {
-    std::map<ggml_backend_t, std::vector<TensorState*>> states_by_compute_backend;
+    std::map<std::pair<ggml_backend_t, ggml_backend_buffer_type_t>, std::vector<TensorState*>> states_by_staging_target;
     for (TensorState* state : states) {
         if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
             continue;
@@ -257,11 +287,16 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
             LOG_ERROR("model manager tensor '%s' is not loaded to params backend", state->name.c_str());
             return false;
         }
-        states_by_compute_backend[state->compute_backend].push_back(state);
+        ggml_backend_buffer_type_t staging_buft = split_buffer_type_for(*state);
+        if (staging_buft == nullptr) {
+            staging_buft = ggml_backend_get_default_buffer_type(state->compute_backend);
+        }
+        states_by_staging_target[{state->compute_backend, staging_buft}].push_back(state);
     }
 
-    for (const auto& pair : states_by_compute_backend) {
-        ggml_backend_t compute_backend          = pair.first;
+    for (const auto& pair : states_by_staging_target) {
+        ggml_backend_t compute_backend          = pair.first.first;
+        ggml_backend_buffer_type_t staging_buft = pair.first.second;
         const std::vector<TensorState*>& states = pair.second;
         if (states.empty()) {
             continue;
@@ -285,7 +320,7 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
             staged_tensors.push_back({state, staging_tensor});
         }
 
-        ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors(staging_ctx, compute_backend);
+        ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, staging_buft);
         if (compute_buffer == nullptr) {
             LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu",
                       staged_tensors.size());
@@ -350,6 +385,17 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
             LOG_ERROR("model manager compute backend is null for lora target tensor '%s'", state->name.c_str());
             return false;
         }
+        if (state->tensor->buffer != nullptr &&
+            ggml_backend_buffer_get_type(state->tensor->buffer) == split_buffer_type_for(*state)) {
+            if (!warned_split_lora_skip_) {
+                LOG_WARN(
+                    "model manager skipping direct lora application to row-split tensors "
+                    "(use --lora-apply-mode at_runtime with row split)");
+                warned_split_lora_skip_ = true;
+            }
+            state->applied_lora_epoch = current_lora_epoch_;
+            continue;
+        }
         if (state->tensor->data == nullptr) {
             LOG_ERROR("model manager lora target tensor '%s' is not prepared", state->name.c_str());
             return false;
@@ -694,6 +740,8 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat
         if (compute_dev != nullptr) {
             params_buft = ggml_backend_dev_host_buffer_type(compute_dev);
         }
+    } else if (state.params_backend == state.compute_backend) {
+        params_buft = split_buffer_type_for(state);
     }
     if (params_buft == nullptr) {
         params_buft = ggml_backend_get_default_buffer_type(state.params_backend);