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

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 @@ 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,

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);

src/model_manager.h

@@ -36,6 +36,7 @@ private:
         ResidencyMode residency_mode   = ResidencyMode::ParamBackend;
         ggml_backend_t compute_backend = nullptr;
         ggml_backend_t params_backend  = nullptr;
+        bool allow_split_buffer        = false;
         bool metadata_validated        = false;
 
         int active_prepare_count = 0;
@@ -63,6 +64,8 @@ private:
     std::map<std::string, TensorState*> tensor_states_by_name_;
     std::vector<std::unique_ptr<ParamsStorageBlock>> params_storage_blocks_;
     std::vector<std::unique_ptr<ComputeStagingBlock>> compute_staging_blocks_;
+    std::map<ggml_backend_t, ggml_backend_buffer_type_t> split_buffer_types_;
+    bool warned_split_lora_skip_ = false;
     std::set<std::string> common_ignore_tensors_;
     std::vector<LoraSpec> loras_;
     SDVersion lora_version_      = VERSION_COUNT;
@@ -91,6 +94,7 @@ private:
     bool stage_tensors_to_compute_backend(const std::vector<TensorState*>& states);
 
     ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const;
+    ggml_backend_buffer_type_t split_buffer_type_for(const TensorState& state) const;
     void release_compute_staging_blocks(bool force                                            = false,
                                         const std::unordered_set<TensorState*>* target_states = nullptr);
     void release_params_storage_blocks(bool force                                            = false,
@@ -114,6 +118,9 @@ public:
     void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; }
     void set_common_ignore_tensors(std::set<std::string> ignore_tensors);
     void set_loras(std::vector<LoraSpec> loras, SDVersion version);
+    void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft);
+
+    static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor);
 
     std::set<std::string> tensor_names() const;
 
@@ -122,7 +129,8 @@ public:
                                 ResidencyMode residency_mode,
                                 ggml_backend_t compute_backend,
                                 ggml_backend_t params_backend,
-                                size_t* registered_tensor_size = nullptr);
+                                size_t* registered_tensor_size = nullptr,
+                                bool allow_split_buffer        = false);
 
     template <typename Runner>
     bool register_runner_params(const std::string& desc,

src/stable-diffusion.cpp

@@ -218,6 +218,7 @@ public:
     bool eager_load    = false;
     std::string backend_spec;
     std::string params_backend_spec;
+    std::string split_mode_spec;
 
     bool is_using_v_parameterization     = false;
     bool is_using_edm_v_parameterization = false;
@@ -292,6 +293,15 @@ public:
         if (module_backends.size() > 1) {
             if constexpr (has_set_runtime_backends<T>::value) {
                 if (module == SDBackendModule::DIFFUSION || module == SDBackendModule::TE) {
+                    if (backend_manager.split_mode(module) == SDSplitMode::ROW) {
+                        return register_row_split_runner_params(desc,
+                                                                model,
+                                                                module,
+                                                                module_backends,
+                                                                std::move(group_tensors),
+                                                                residency_mode,
+                                                                params_mem_size);
+                    }
                     return register_layer_split_runner_params(desc,
                                                               model,
                                                               module,
@@ -313,6 +323,108 @@ public:
                                                      params_mem_size);
     }
 
+    template <typename T>
+    bool register_row_split_runner_params(const std::string& desc,
+                                          const std::shared_ptr<T>& model,
+                                          SDBackendModule module,
+                                          const std::vector<ggml_backend_t>& module_backends,
+                                          std::map<std::string, ggml_tensor*> group_tensors,
+                                          ModelManager::ResidencyMode residency_mode,
+                                          size_t* params_mem_size) {
+        ggml_backend_t main_backend = module_backends[0];
+
+        auto fall_back_to_layer_split = [&](const char* reason) {
+            LOG_WARN("%s: row split unavailable (%s); falling back to layer split", desc.c_str(), reason);
+            return register_layer_split_runner_params(desc,
+                                                      model,
+                                                      module,
+                                                      module_backends,
+                                                      std::move(group_tensors),
+                                                      residency_mode,
+                                                      params_mem_size);
+        };
+
+        ggml_backend_dev_t main_dev = ggml_backend_get_device(main_backend);
+        ggml_backend_reg_t reg      = main_dev != nullptr ? ggml_backend_dev_backend_reg(main_dev) : nullptr;
+        if (reg == nullptr) {
+            return fall_back_to_layer_split("no backend registry");
+        }
+        const size_t reg_dev_count = ggml_backend_reg_dev_count(reg);
+        std::vector<float> tensor_split(reg_dev_count, 0.0f);
+        constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
+        for (ggml_backend_t backend : module_backends) {
+            ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+            int reg_index          = -1;
+            for (size_t i = 0; i < reg_dev_count; i++) {
+                if (ggml_backend_reg_dev_get(reg, i) == dev) {
+                    reg_index = (int)i;
+                    break;
+                }
+            }
+            if (reg_index < 0) {
+                return fall_back_to_layer_split("devices span different backend registries");
+            }
+            size_t free_bytes = 0, total_bytes = 0;
+            ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
+            int64_t usable_bytes    = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes,
+                                                     (int64_t)free_bytes / 8);
+            tensor_split[reg_index] = usable_bytes > 0 ? (float)((double)usable_bytes / (1024.0 * 1024.0)) : 1.0f;
+        }
+
+        ggml_backend_buffer_type_t split_buft = backend_manager.split_buffer_type(main_backend, tensor_split);
+        if (split_buft == nullptr) {
+            return fall_back_to_layer_split("backend has no split buffer type");
+        }
+        model_manager->set_split_buffer_type(main_backend, split_buft);
+
+        std::map<std::string, ggml_tensor*> split_tensors;
+        if constexpr (std::is_base_of_v<Conditioner, T>) {
+            model->get_layer_split_param_tensors(split_tensors);
+        } else {
+            split_tensors = group_tensors;
+        }
+
+        std::map<std::string, ggml_tensor*> row_split_map;
+        std::map<std::string, ggml_tensor*> regular_map;
+        size_t row_split_bytes = 0;
+        for (const auto& kv : group_tensors) {
+            if (split_tensors.count(kv.first) != 0 &&
+                sd::layer_split_tensor_block_index(kv.first) >= 0 &&
+                ModelManager::tensor_shape_supports_split_buffer(kv.second)) {
+                row_split_map[kv.first] = kv.second;
+                row_split_bytes += ggml_nbytes(kv.second);
+            } else {
+                regular_map[kv.first] = kv.second;
+            }
+        }
+        if (row_split_map.empty()) {
+            return fall_back_to_layer_split("no row-splittable transformer block weights found");
+        }
+
+        LOG_INFO("%s row split: %zu tensors (%.1f MB) split across %zu devices (main %s)",
+                 desc.c_str(),
+                 row_split_map.size(),
+                 row_split_bytes / (1024.f * 1024.f),
+                 module_backends.size(),
+                 sd::layer_split_backend_device_display_name(main_backend).c_str());
+
+        if (!model_manager->register_param_tensors(desc,
+                                                   std::move(row_split_map),
+                                                   residency_mode,
+                                                   main_backend,
+                                                   params_backend_for(module),
+                                                   params_mem_size,
+                                                   /*allow_split_buffer=*/true)) {
+            return false;
+        }
+        return model_manager->register_param_tensors(desc,
+                                                     std::move(regular_map),
+                                                     residency_mode,
+                                                     main_backend,
+                                                     params_backend_for(module),
+                                                     params_mem_size);
+    }
+
     // Register each layer-split partition with its compute backend; the
     // ModelManager handles allocation, staging, and LoRA by backend.
     template <typename T>
@@ -397,6 +509,7 @@ public:
         std::string error;
         if (!backend_manager.init(backend_spec.c_str(),
                                   params_backend_spec.c_str(),
+                                  split_mode_spec.c_str(),
                                   &error)) {
             LOG_ERROR("backend config failed: %s", error.c_str());
             return false;
@@ -404,6 +517,16 @@ public:
         return ensure_backend_pair(SDBackendModule::DIFFUSION);
     }
 
+    bool row_split_active() {
+        for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) {
+            if (backend_manager.split_mode(module) == SDSplitMode::ROW &&
+                backend_manager.runtime_backends(module).size() > 1) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     std::shared_ptr<RNG> get_rng(rng_type_t rng_type) {
         if (rng_type == STD_DEFAULT_RNG) {
             return std::make_shared<STDDefaultRNG>();
@@ -462,6 +585,7 @@ public:
         eager_load          = sd_ctx_params->eager_load;
         backend_spec        = SAFE_STR(sd_ctx_params->backend);
         params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
+        split_mode_spec     = SAFE_STR(sd_ctx_params->split_mode);
         max_vram_assignment.reset(0.f);
         {
             std::string error;
@@ -690,12 +814,17 @@ public:
             // Avoid full-model LoRA merge buffers on constrained setups.
             const bool params_offloaded      = params_backend_for(SDBackendModule::DIFFUSION) != backend_for(SDBackendModule::DIFFUSION);
             const bool streaming_constrained = stream_layers || params_offloaded;
-            if (have_quantized_weight || streaming_constrained) {
+            if (have_quantized_weight || streaming_constrained || row_split_active()) {
                 apply_lora_immediately = false;
             } else {
                 apply_lora_immediately = true;
             }
         } else if (sd_ctx_params->lora_apply_mode == LORA_APPLY_IMMEDIATELY) {
+            if (row_split_active()) {
+                LOG_WARN(
+                    "row-split tensors do not support the immediately LoRA apply mode; "
+                    "LoRAs will not be applied to them (use --lora-apply-mode at_runtime)");
+            }
             apply_lora_immediately = true;
         } else {
             apply_lora_immediately = false;
@@ -2916,6 +3045,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_params->vae_format           = SD_VAE_FORMAT_AUTO;
     sd_ctx_params->backend              = nullptr;
     sd_ctx_params->params_backend       = nullptr;
+    sd_ctx_params->split_mode           = nullptr;
     sd_ctx_params->rpc_servers          = nullptr;
     sd_ctx_params->pulid_weights_path   = nullptr;
 }
@@ -2955,6 +3085,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "eager_load: %s\n"
              "backend: %s\n"
              "params_backend: %s\n"
+             "split_mode: %s\n"
              "flash_attn: %s\n"
              "diffusion_flash_attn: %s\n"
              "circular_x: %s\n"
@@ -2991,6 +3122,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              BOOL_STR(sd_ctx_params->eager_load),
              SAFE_STR(sd_ctx_params->backend),
              SAFE_STR(sd_ctx_params->params_backend),
+             SAFE_STR(sd_ctx_params->split_mode),
              BOOL_STR(sd_ctx_params->flash_attn),
              BOOL_STR(sd_ctx_params->diffusion_flash_attn),
              BOOL_STR(sd_ctx_params->circular_x),

src/upscaler.cpp

@@ -46,6 +46,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
     std::string error;
     if (!backend_manager.init(backend_spec.c_str(),
                               params_backend_spec.c_str(),
+                              /*split_mode_spec=*/nullptr,
                               &error)) {
         LOG_ERROR("upscaler backend config failed: %s", error.c_str());
         return false;