diff --git a/src/llama-model.h b/src/llama-model.h index c9bae3103..73447913e 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -493,6 +493,11 @@ struct llama_model { std::unique_ptr output_mtp_ptr; + // Device-local DFlash IO copies for cross-buffer sharing. + std::unique_ptr dflash_tok_embd_ptr; + std::unique_ptr dflash_output_ptr; + std::unique_ptr dflash_output_mtp_ptr; + llama_split_tensor split_output; llama_split_tensor split_output_norm; diff --git a/src/llama-spec-features-dflash.cpp b/src/llama-spec-features-dflash.cpp index 01ad482c9..7348388a6 100644 --- a/src/llama-spec-features-dflash.cpp +++ b/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 & storage, + const char * name) { + if (model == nullptr || source == nullptr || source->buffer == nullptr || buft == nullptr) { + return nullptr; + } + + storage = std::make_unique(*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; }