Prefil optimizations for draft model

common/speculative.cpp

@@ -1103,29 +1103,78 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
             for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) {
                 const int32_t n_chunk = std::min(n_ubatch, n_rows - offset);
 
-                // gather this chunk's target features, interleaved by extract layer
-                features_buf.resize((size_t) n_chunk * n_embd_enc);
-                for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
-                    const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
-                    if (!layer) {
-                        GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]);
+                // fuse extracted features through DFlash encoder.
+                //
+                // zero-copy path (preferred): the target context concatenates the enabled
+                // layer-input tensors into a persistent device buffer (embd_layer_inp_fused)
+                // each compute call; we alias it via embd_dev (a view, no host copy).
+                // falls back to the host gather below when the device/buffer types are not
+                // compatible (e.g. target and draft on different GPUs).
+                // zero-copy requires the fused tensor to live on a backend that the draft
+                // encoder can consume directly (same device, compatible buffer type).
+                // otherwise fall back to the host path.
+                //
+                // when no explicit devices are configured, the draft and target share the
+                // default devices (single-GPU case) - assume compatible.
+                const ggml_tensor * fused = llama_get_embeddings_layer_inp_tensor(ctx_tgt);
+                if (fused && fused->buffer && !this->params.devices.empty()) {
+                    const ggml_backend_dev_t fused_dev = ggml_backend_buft_get_device(
+                            ggml_backend_buffer_get_type(fused->buffer));
+                    bool compatible = false;
+                    for (const auto & dev : this->params.devices) {
+                        if (dev == fused_dev) {
+                            compatible = true;
+                            break;
+                        }
+                    }
+                    if (!compatible) {
+                        fused = nullptr;
                     }
-                    for (int32_t i = 0; i < n_chunk; ++i) {
-                        float       * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
-                        const float * src = layer + (size_t) (i_batch_beg[seq_id] + offset + i) * n_embd_tgt;
-                        std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));
+                }
+                if (fused) {
+                    // event-based cross-stream sync: the draft backends wait on the GPU
+                    // stream for the fused write to complete (no host block).
+                    llama_embd_layer_inp_wait(ctx_tgt, ctx_dft);
+                }
+                static bool fused_logged = false;
+                if (!fused_logged) {
+                    fused_logged = true;
+                    if (fused) {
+                        LOG_INF("%s: DFlash zero-copy embd path active (fused tensor %s)\n",
+                                __func__, fused->name ? fused->name : "?");
+                    } else {
+                        LOG_INF("%s: DFlash host embd path active (no fused tensor)\n", __func__);
+                    }
+                }
+                // NOTE: the fused buffer is overwritten by the next llama_decode(ctx_tgt).
+                //       the speculative loop guarantees the draft consumes it before then:
+                //       draft() -> llama_decode(ctx_tgt) -> process() -> verify -> repeat.
+                if (!fused) {
+                    // host path: gather this chunk's target features, interleaved by extract layer
+                    features_buf.resize((size_t) n_chunk * n_embd_enc);
+                    for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
+                        const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
+                        if (!layer) {
+                            GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]);
+                        }
+                        for (int32_t i = 0; i < n_chunk; ++i) {
+                            float       * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
+                            const float * src = layer + (size_t) (i_batch_beg[seq_id] + offset + i) * n_embd_tgt;
+                            std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));
+                        }
                     }
                 }
 
-                // fuse extracted features through DFlash encoder
                 llama_batch enc_batch = {
                     /*.n_tokens =*/ n_chunk,
                     /*.token    =*/ nullptr,
-                    /*.embd     =*/ features_buf.data(),
+                    /*.embd     =*/ fused ? nullptr : features_buf.data(),
                     /*.pos      =*/ nullptr,
                     /*.n_seq_id =*/ nullptr,
                     /*.seq_id   =*/ nullptr,
                     /*.logits   =*/ nullptr,
+                    /*.embd_dev =*/ (ggml_tensor *) fused,
+                    /*.embd_dev_off =*/ (int64_t) i_batch_beg[seq_id] + offset,
                 };
 
                 int32_t rc = llama_encode(ctx_dft, enc_batch);
@@ -1135,12 +1184,24 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                     return false;
                 }
 
-                const float * inp_g = llama_get_embeddings_nextn(ctx_dft);
-                GGML_ASSERT(inp_g && "DFlash encoder produced no output.");
+                // zero-copy nextn path: the encoder wrote its output (t_h_nextn) into a
+                // persistent device buffer; alias it for the decoder KV-injection instead
+                // of a host read + H2D copy. same context/stream, so ordering is guaranteed.
+                const ggml_tensor * nextn_persist = llama_get_embeddings_nextn_tensor(ctx_dft);
+                if (nextn_persist) {
+                    batch_inject.embd     = nullptr;
+                    batch_inject.embd_dev = (ggml_tensor *) nextn_persist;
+                } else {
+                    const float * inp_g = llama_get_embeddings_nextn(ctx_dft);
+                    GGML_ASSERT(inp_g && "DFlash encoder produced no output.");
+
+                    batch_inject.embd_dev = nullptr;
+                    std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float));
+                }
 
                 // inject the DFlash decoder K/V cache at the tokens' target positions
                 batch_inject.n_tokens = n_chunk;
-                std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float));
+                batch_inject.embd_dev_off = 0;
 
                 for (int32_t i = 0; i < n_chunk; ++i) {
                     batch_inject.pos[i]       = batch_in.pos[i_batch_beg[seq_id] + offset + i];

ggml/src/ggml-impl.h

@@ -11,6 +11,9 @@
 #include <stdbool.h>
 #include <stdint.h>
 #include <string.h>
+#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__)
+    #include <immintrin.h>
+#endif
 
 #ifdef __ARM_FEATURE_SVE
 #include <arm_sve.h>
@@ -382,6 +385,12 @@ static inline uint32_t fp32_to_bits(float f) {
 }
 
 static inline float ggml_compute_fp16_to_fp32(ggml_fp16_t h) {
+#ifdef __F16C__
+    return _cvtsh_ss(h);
+#elif defined(__aarch64__) && defined(__ARM_FP) && (__ARM_FP & 2)
+    union { uint16_t u; __fp16 f; } u = { .u = h };
+    return (float)u.f;
+#else
     const uint32_t w = (uint32_t) h << 16;
     const uint32_t sign = w & UINT32_C(0x80000000);
     const uint32_t two_w = w + w;
@@ -402,6 +411,7 @@ static inline float ggml_compute_fp16_to_fp32(ggml_fp16_t h) {
     const uint32_t result = sign |
         (two_w < denormalized_cutoff ? fp32_to_bits(denormalized_value) : fp32_to_bits(normalized_value));
     return fp32_from_bits(result);
+#endif
 }
 
 static inline ggml_fp16_t ggml_compute_fp32_to_fp16(float f) {

include/llama.h

@@ -262,6 +262,13 @@ extern "C" {
         int32_t      *  n_seq_id;
         llama_seq_id ** seq_id;
         int8_t       *  logits;   // TODO: rename this to "output"
+
+        // device-side embedding input (zero-copy path)
+        // when set, `embd` is ignored: the graph consumes a view of this tensor
+        // (owned by another context, e.g. the speculative target) starting at
+        // row `embd_dev_off`. both fields are optional and ignored by the host path.
+        struct ggml_tensor * embd_dev;
+        int64_t              embd_dev_off;
     } llama_batch;
 
     enum llama_model_kv_override_type {

src/llama-batch.cpp

@@ -218,6 +218,8 @@ bool llama_batch_allocr::init(
             /*.n_pos        =*/ n_pos_per_embd,
             /*.token        =*/ batch.token,
             /*.embd         =*/ batch.embd,
+            /*.embd_dev     =*/ batch.embd_dev,
+            /*.embd_dev_off =*/ batch.embd_dev_off,
             /*.pos          =*/ batch.pos,
             /*.n_seq_id     =*/ batch.n_seq_id,
             /*.seq_id       =*/ batch.seq_id,
@@ -424,6 +426,8 @@ llama_ubatch llama_batch_allocr::ubatch_reserve(uint32_t n_seq_tokens, uint32_t
 
         /*.token        =*/ udata->token.data(),
         /*.embd         =*/ nullptr,
+        /*.embd_dev     =*/ nullptr,
+        /*.embd_dev_off =*/ 0,
         /*.pos          =*/ udata->pos.data(),
         /*.n_seq_id     =*/ udata->n_seq_id.data(),
         /*.seq_id       =*/ udata->seq_id.data(),
@@ -753,11 +757,12 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
 
     auto udata = std::make_shared<llama_ubatch::data_t>();
 
-    const int64_t n_embd_all = batch.embd ? (int64_t) n_tokens*n_embd : 0;
-    const int64_t n_pos_all  =              (int64_t) n_tokens*n_pos_per_embd;
+    // host embedding buffer is needed only for the host path (embd set, no device alias)
+    const bool has_embd_host = batch.embd && !batch.embd_dev;
+    const int64_t n_embd_all = has_embd_host ? (int64_t) n_tokens*n_embd : 0;
+    const int64_t n_pos_all  =                 (int64_t) n_tokens*n_pos_per_embd;
 
     udata->token     .resize(n_tokens);
-    udata->embd      .resize(n_embd_all);
     udata->pos       .resize(n_pos_all);
     udata->n_seq_id  .resize(n_tokens);
     udata->seq_id    .resize(n_tokens);
@@ -766,6 +771,21 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
     udata->output    .resize(n_tokens);
 
     udata->seq_id_data.reserve(n_tokens);
+    if (has_embd_host) {
+        udata->embd.clear();
+        udata->embd.reserve(n_embd_all);
+    } else {
+        udata->embd.resize(n_embd_all); // fill all size..new_size elems by 0.0f
+    }
+
+    if (batch.embd_dev) {
+        // zero-copy device alias: the external tensor is a single contiguous
+        // [n_embd, n_tokens] block; we can only alias a contiguous run of indices
+        // (the current callers - e.g. the dflash encoder ubatch - provide one)
+        for (size_t i = 0; i < idxs.size(); ++i) {
+            GGML_ASSERT(idxs[i] == idxs[0] + (int32_t) i);
+        }
+    }
 
     seq_set_t seq_set_unq;
 
@@ -774,8 +794,11 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
             udata->token[i] = batch.token[idxs[i]];
         }
 
-        if (batch.embd) {
-            memcpy(udata->embd.data() + i*n_embd, batch.embd + (int64_t) idxs[i]*n_embd, n_embd*sizeof(float));
+        if (has_embd_host) {
+            auto src = batch.embd + (int64_t) idxs[i] * n_embd;
+            // use safe method for auto increase size
+            // next improvements - write own vector without automatic filling float)
+            udata->embd.insert(udata->embd.end(), src, src + n_embd);
         }
 
         for (size_t j = 0; j < (size_t)n_pos_per_embd; ++j) {
@@ -824,7 +847,9 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
         /*.n_pos        =*/ n_pos_per_embd,
 
         /*.token        =*/ batch.token ? udata->token.data() : nullptr,
-        /*.embd         =*/ batch.embd ? udata->embd.data() : nullptr,
+        /*.embd         =*/ (batch.embd && !batch.embd_dev) ? udata->embd.data() : nullptr,
+        /*.embd_dev     =*/ batch.embd_dev,
+        /*.embd_dev_off =*/ batch.embd_dev_off + idxs[0],
         /*.pos          =*/ udata->pos.data(),
         /*.n_seq_id     =*/ udata->n_seq_id.data(),
         /*.seq_id       =*/ udata->seq_id.data(),
@@ -874,6 +899,7 @@ void llama_batch_allocr::ubatch_print(const llama_ubatch & ubatch, int debug) {
 
         LLAMA_LOG_DEBUG("%s:   token      = %p\n", __func__, (void *) ubatch.token);
         LLAMA_LOG_DEBUG("%s:   embd       = %p\n", __func__, (void *) ubatch.embd);
+        LLAMA_LOG_DEBUG("%s:   embd_dev   = %p (off = %ld)\n", __func__, (void *) ubatch.embd_dev, (long) ubatch.embd_dev_off);
         LLAMA_LOG_DEBUG("%s:   pos        = %p\n", __func__, (void *) ubatch.pos);
         LLAMA_LOG_DEBUG("%s:   n_seq_id   = %p\n", __func__, (void *) ubatch.n_seq_id);
         LLAMA_LOG_DEBUG("%s:   seq_id     = %p\n", __func__, (void *) ubatch.seq_id);
@@ -939,6 +965,8 @@ struct llama_batch llama_batch_get_one(
         /*n_seq_id =*/ nullptr,
         /*seq_id   =*/ nullptr,
         /*logits   =*/ nullptr,
+        /*embd_dev =*/ nullptr,
+        /*embd_dev_off =*/ 0,
     };
 }
 
@@ -951,6 +979,8 @@ struct llama_batch llama_batch_init(int32_t n_tokens_alloc, int32_t embd, int32_
         /*n_seq_id =*/ nullptr,
         /*seq_id   =*/ nullptr,
         /*logits   =*/ nullptr,
+        /*embd_dev =*/ nullptr,
+        /*embd_dev_off =*/ 0,
     };
 
     if (embd) {

src/llama-batch.h

@@ -44,6 +44,14 @@ struct llama_ubatch {
     //                          // size               | idx | val
     llama_token  *  token;      // [n_tokens]         | i   | id, token
     float        *  embd;       // [n_embd, n_tokens] | i   | embd
+
+    // device-side embedding input (zero-copy path)
+    // when set, `embd` is ignored: the graph consumes a view of this external
+    // tensor (owned by another context, e.g. the speculative target) starting
+    // at column `embd_dev_off`. both fields are optional and ignored by the host path.
+    struct ggml_tensor * embd_dev;
+    int64_t              embd_dev_off;
+
     llama_pos    *  pos;        // [n_tokens*n_pos]   | i   | pos
     int32_t      *  n_seq_id;   // [n_tokens]         | i   | -
     llama_seq_id ** seq_id;     // [n_tokens]         | s   | s0, s1, seq_id