fix: isolate DFlash cross-device IO (#2243)

* fix: isolate DFlash cross-device IO

* remove unnecesary syncronization and simplify logic for tensors copies

src/llama-model.h

@@ -493,6 +493,11 @@ struct llama_model {
 
     std::unique_ptr<ggml_tensor> output_mtp_ptr;
 
+    // Device-local DFlash IO copies for cross-buffer sharing.
+    std::unique_ptr<ggml_tensor> dflash_tok_embd_ptr;
+    std::unique_ptr<ggml_tensor> dflash_output_ptr;
+    std::unique_ptr<ggml_tensor> dflash_output_mtp_ptr;
+
     llama_split_tensor split_output;
     llama_split_tensor split_output_norm;
 

src/llama-spec-features-dflash.cpp

@@ -158,6 +158,45 @@ int32_t llama_model_dflash_io_mode(
     return LLAMA_DFLASH_IO_MODE_MIXED;
 }
 
+static ggml_tensor * llama_dflash_clone_io_tensor(
+        llama_model * model,
+        ggml_tensor * source,
+        ggml_backend_buffer_type_t buft,
+        std::unique_ptr<ggml_tensor> & storage,
+        const char * name) {
+    if (model == nullptr || source == nullptr || source->buffer == nullptr || buft == nullptr) {
+        return nullptr;
+    }
+
+    storage = std::make_unique<ggml_tensor>(*source);
+    storage->buffer = ggml_backend_buft_alloc_buffer(buft, ggml_backend_buft_get_alloc_size(buft, source));
+    if (storage->buffer == nullptr) {
+        storage.reset();
+        return nullptr;
+    }
+
+    storage->data = ggml_backend_buffer_get_base(storage->buffer);
+    storage->op = GGML_OP_NONE;
+    for (int j = 0; j < GGML_MAX_SRC; ++j) {
+        storage->src[j] = nullptr;
+    }
+    storage->view_src = nullptr;
+    storage->view_offs = 0;
+    storage->extra = nullptr;
+    ggml_set_name(storage.get(), name);
+    ggml_backend_buffer_set_usage(storage->buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+
+    ggml_backend_tensor_copy(source, storage.get());
+
+    model->bufs.push_back(storage->buffer);
+    return storage.get();
+}
+
+static bool llama_dflash_io_needs_clone(const ggml_tensor * tensor, ggml_backend_buffer_type_t buft) {
+    return tensor != nullptr && tensor->buffer != nullptr && buft != nullptr &&
+           ggml_backend_buffer_get_type(tensor->buffer) != buft;
+}
+
 bool llama_model_dflash_io_tensors_match(
         const struct llama_model * draft_model,
         int32_t n_embd,
@@ -209,6 +248,52 @@ bool llama_model_share_dflash_io_tensors(
         }
     }
 
+    const bool output_mtp_aliases_output = draft_model->output_mtp == draft_model->output;
+    const bool tok_embd_is_shared = draft_model->tok_embd == target_model->tok_embd;
+    const bool output_is_shared = draft_model->output == target_model->output ||
+            draft_model->output == target_model->tok_embd;
+    const bool output_mtp_is_shared = draft_model->output_mtp == target_model->output_mtp ||
+            draft_model->output_mtp == target_model->output ||
+            draft_model->output_mtp == target_model->tok_embd;
+    const bool isolate_shared_io =
+            (tok_embd_is_shared && llama_dflash_io_needs_clone(draft_model->tok_embd, draft_model->buft_input.buft)) ||
+            (output_is_shared && llama_dflash_io_needs_clone(draft_model->output, draft_model->buft_output.buft)) ||
+            (output_mtp_is_shared && llama_dflash_io_needs_clone(draft_model->output_mtp, draft_model->buft_output.buft));
+
+    if (tok_embd_is_shared && (isolate_shared_io || llama_dflash_io_needs_clone(draft_model->tok_embd, draft_model->buft_input.buft))) {
+        ggml_tensor * source = draft_model->tok_embd;
+        draft_model->tok_embd = llama_dflash_clone_io_tensor(
+                draft_model, source, draft_model->buft_input.buft, draft_model->dflash_tok_embd_ptr,
+                "dflash_tok_embd");
+        if (draft_model->tok_embd == nullptr) {
+            return false;
+        }
+    }
+
+    if (output_is_shared && (isolate_shared_io || llama_dflash_io_needs_clone(draft_model->output, draft_model->buft_output.buft))) {
+        ggml_tensor * source = draft_model->output;
+        draft_model->output = llama_dflash_clone_io_tensor(
+                draft_model, source, draft_model->buft_output.buft, draft_model->dflash_output_ptr,
+                "dflash_output");
+        if (draft_model->output == nullptr) {
+            return false;
+        }
+        if (output_mtp_aliases_output) {
+            draft_model->output_mtp = draft_model->output;
+        }
+    }
+
+    if (!output_mtp_aliases_output && output_mtp_is_shared &&
+            (isolate_shared_io || llama_dflash_io_needs_clone(draft_model->output_mtp, draft_model->buft_output.buft))) {
+        ggml_tensor * source = draft_model->output_mtp;
+        draft_model->output_mtp = llama_dflash_clone_io_tensor(
+                draft_model, source, draft_model->buft_output.buft, draft_model->dflash_output_mtp_ptr,
+                "dflash_output_mtp");
+        if (draft_model->output_mtp == nullptr) {
+            return false;
+        }
+    }
+
     const struct ggml_tensor * output = llama_dflash_output_tensor(draft_model);
     return draft_model->tok_embd != nullptr && output != nullptr;
 }