diff --git a/common/speculative.cpp b/common/speculative.cpp index ae55e357d51a..5bddb1985b4d 100644 --- a/common/speculative.cpp +++ b/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]; diff --git a/ggml/src/ggml-impl.h b/ggml/src/ggml-impl.h index 62b76abbcec9..7189f9539f07 100644 --- a/ggml/src/ggml-impl.h +++ b/ggml/src/ggml-impl.h @@ -11,6 +11,9 @@ #include #include #include +#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) + #include +#endif #ifdef __ARM_FEATURE_SVE #include @@ -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) { diff --git a/include/llama.h b/include/llama.h index 177fc10a9139..bb7a1b747856 100644 --- a/include/llama.h +++ b/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 { diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2b98a552f48f..ea2b3c94c9cf 100644 --- a/src/llama-batch.cpp +++ b/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 & idxs, u auto udata = std::make_shared(); - 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 & 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 & 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 & 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) { diff --git a/src/llama-batch.h b/src/llama-batch.h index a3d1889d4a04..83475f9c3686 100644 --- a/src/llama-batch.h +++ b/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 diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 52f8d53672a3..ba3d2d44111e 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -482,6 +482,11 @@ llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); + if (embd_layer_inp_fused_event) { + ggml_backend_event_free(embd_layer_inp_fused_event); + embd_layer_inp_fused_event = nullptr; + } + if (!model.hparams.no_alloc) { for (size_t i = 0; i < backend_ptrs.size(); ++i) { ggml_backend_t backend = backend_ptrs[i]; @@ -1322,7 +1327,7 @@ bool llama_context::set_adapter_cvec( return res; } -llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, llm_graph_type gtype, llama_memory_context_i * mctx, ggml_status & ret) { +llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, llm_graph_type gtype, llama_memory_context_i * mctx, ggml_status & ret, int64_t token_offset) { if (mctx && !mctx->apply()) { LLAMA_LOG_ERROR("%s: failed to apply memory context\n", __func__); ret = GGML_STATUS_FAILED; @@ -1332,9 +1337,14 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll auto * res = gf_res_prev.get(); auto * gf = res->get_gf(); + // ensure the persistent fused layer-input device buffer exists (zero-copy path), + // so that graph_params can hand it to the graph builder + ensure_embd_layer_inp_fused(); + ensure_embd_nextn_persist(); + // the new graph parameters // in order to correctly reuse a graph, it's full topology has to be uniquely determined by these parameters - const auto gparams = graph_params(res, ubatch, mctx, gtype); + const auto gparams = graph_params(res, ubatch, mctx, gtype, token_offset); if (!graph_reuse_disable && res->can_reuse(gparams)) { //LLAMA_LOG_DEBUG("%s: reusing previous graph\n", __func__); @@ -1389,6 +1399,16 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll return nullptr; } + // record the fused layer-input write event so a foreign context (the speculative + // draft) can wait on the GPU stream instead of host-syncing. graph_compute runs + // async, so the event lands in the queue after the fused cpy kernels. + if (embd_layer_inp_fused_event && embd_layer_inp_fused) { + ggml_backend_t fused_backend = ggml_backend_sched_get_tensor_backend(sched.get(), embd_layer_inp_fused); + if (fused_backend) { + ggml_backend_event_record(embd_layer_inp_fused_event, fused_backend); + } + } + ret = GGML_STATUS_SUCCESS; return res; @@ -1397,7 +1417,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll int llama_context::encode(const llama_batch & batch_inp) { // MTP hook batches carry both token (next-token id) and embd (h_nextn row), // so accept either present rather than requiring exactly one. - GGML_ASSERT(batch_inp.token || batch_inp.embd); + GGML_ASSERT(batch_inp.token || batch_inp.embd || batch_inp.embd_dev); if (batch_inp.n_tokens == 0) { LLAMA_LOG_ERROR("%s: n_tokens == 0\n", __func__); @@ -1635,7 +1655,7 @@ static bool needs_raw_logits(const llama_ubatch & ubatch, const std::map remove all positions of that ubatch from the memory module @@ -2219,6 +2239,106 @@ void llama_context::extract_layer_inputs(const llm_graph_result * res, size_t to } } +ggml_tensor * llama_context::ensure_embd_layer_inp_fused() { + // count enabled layer-input extractions; without any, the zero-copy path is disabled + uint32_t n_extract = 0; + for (bool enabled : cparams.embeddings_layer_inp) { + if (enabled) { + ++n_extract; + } + } + if (n_extract == 0) { + return nullptr; + } + + if (embd_layer_inp_fused) { + return embd_layer_inp_fused; + } + + const uint32_t n_embd = model.hparams.n_embd; + const uint32_t n_batch = cparams.n_batch; + + auto * dev = model.dev_output(); + if (dev == nullptr) { + return nullptr; + } + auto * buft = ggml_backend_dev_buffer_type(dev); + + ggml_init_params iparams = { + /*.mem_size =*/ ggml_tensor_overhead() + ggml_graph_overhead_custom(0, false), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + embd_layer_inp_fused_ctx = ggml_init(iparams); + if (!embd_layer_inp_fused_ctx) { + return nullptr; + } + + embd_layer_inp_fused = ggml_new_tensor_2d(embd_layer_inp_fused_ctx, GGML_TYPE_F32, + (int64_t) n_extract * n_embd, n_batch); + ggml_set_name(embd_layer_inp_fused, "embd_layer_inp_fused"); + + embd_layer_inp_fused_buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(embd_layer_inp_fused_ctx, buft)); + if (!embd_layer_inp_fused_buf) { + embd_layer_inp_fused_ctx = nullptr; + embd_layer_inp_fused = nullptr; + return nullptr; + } + ggml_backend_buffer_set_usage(embd_layer_inp_fused_buf.get(), GGML_BACKEND_BUFFER_USAGE_COMPUTE); + + embd_layer_inp_fused_event = ggml_backend_event_new(dev); + if (!embd_layer_inp_fused_event) { + LLAMA_LOG_WARN("%s: failed to create fused event, falling back to host sync\n", __func__); + // not fatal - the getter falls back to ctx->synchronize() when the event is null + } + + return embd_layer_inp_fused; +} + +ggml_tensor * llama_context::ensure_embd_nextn_persist() { + // only meaningful when nextn extraction is enabled (the speculative draft/encoder + // context); otherwise the buffer would be allocated on every context in vain + if (!cparams.embeddings_nextn) { + return nullptr; + } + + if (embd_nextn_persist) { + return embd_nextn_persist; + } + + const uint32_t n_embd_out = model.hparams.n_embd_out(); + const uint32_t n_batch = cparams.n_batch; + + auto * dev = model.dev_output(); + if (dev == nullptr) { + return nullptr; + } + auto * buft = ggml_backend_dev_buffer_type(dev); + + ggml_init_params iparams = { + /*.mem_size =*/ ggml_tensor_overhead() + ggml_graph_overhead_custom(0, false), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + embd_nextn_persist_ctx = ggml_init(iparams); + if (!embd_nextn_persist_ctx) { + return nullptr; + } + + embd_nextn_persist = ggml_new_tensor_2d(embd_nextn_persist_ctx, GGML_TYPE_F32, n_embd_out, n_batch); + ggml_set_name(embd_nextn_persist, "embd_nextn_persist"); + + embd_nextn_persist_buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(embd_nextn_persist_ctx, buft)); + if (!embd_nextn_persist_buf) { + embd_nextn_persist_ctx = nullptr; + embd_nextn_persist = nullptr; + return nullptr; + } + ggml_backend_buffer_set_usage(embd_nextn_persist_buf.get(), GGML_BACKEND_BUFFER_USAGE_COMPUTE); + + return embd_nextn_persist; +} + void llama_context::output_reorder() { const uint64_t n_vocab = model.vocab.n_tokens(); const uint64_t n_embd = model.hparams.n_embd; @@ -2452,7 +2572,8 @@ llm_graph_params llama_context::graph_params( llm_graph_result * res, const llama_ubatch & ubatch, const llama_memory_context_i * mctx, - llm_graph_type gtype) const { + llm_graph_type gtype, + int64_t token_offset) const { return { /*.arch =*/ model.arch, /*.hparams =*/ model.hparams, @@ -2465,6 +2586,9 @@ llm_graph_params llama_context::graph_params( /*.loras =*/ loras.get(), /*.mctx =*/ mctx, /*.cross =*/ &cross, + /*.token_offset =*/ token_offset, + /*.embd_layer_inp_fused =*/ embd_layer_inp_fused, + /*.embd_nextn_persist =*/ embd_nextn_persist, /*.samplers =*/ sampling.samplers, /*.n_outputs =*/ n_outputs, /*.cb =*/ graph_get_cb(), @@ -3812,6 +3936,56 @@ float * llama_get_embeddings_layer_inp(llama_context * ctx, uint32_t lid) { return ctx->get_embeddings_layer_inp(lid); } +const ggml_tensor * llama_get_embeddings_layer_inp_tensor(llama_context * ctx) { + // NOTE: no ctx->synchronize() here - the consumer (draft) must call + // llama_embd_layer_inp_wait() to wait on the GPU stream for the fused write + // to complete (event-based, no host block). if the consumer reads the tensor + // data on the host, it must synchronize itself. + return ctx->get_embd_layer_inp_fused(); +} + +const ggml_tensor * llama_get_embeddings_nextn_tensor(llama_context * ctx) { + // NOTE: no ctx->synchronize() - the encoder output is consumed on the same context's + // stream (decoder KV-injection runs right after encode), so stream ordering suffices. + return ctx->get_embd_nextn_persist(); +} + + +void llama_embd_layer_inp_wait(llama_context * ctx_tgt, llama_context * ctx_dft) { + if (!ctx_tgt || !ctx_dft) { + return; + } + auto * event = ctx_tgt->get_embd_layer_inp_fused_event(); + if (!event) { + // no event (fused not in use, or event creation failed) - host sync as fallback + ctx_tgt->synchronize(); + return; + } + + auto * sched = ctx_dft->get_sched(); + if (!sched) { + ctx_tgt->synchronize(); + return; + } + + const int n_backends = ggml_backend_sched_get_n_backends(sched); + for (int i = 0; i < n_backends; ++i) { + ggml_backend_t backend = ggml_backend_sched_get_backend(sched, i); + if (backend) { + // only device backends implement event_wait (CUDA/Metal/Vulkan); CPU and + // ACCEL backends don't and don't need to wait (fused lives on the GPU) + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + if (dev) { + const enum ggml_backend_dev_type dev_type = ggml_backend_dev_type(dev); + if (dev_type == GGML_BACKEND_DEVICE_TYPE_CPU || dev_type == GGML_BACKEND_DEVICE_TYPE_ACCEL) { + continue; + } + } + ggml_backend_event_wait(backend, event); + } + } +} + bool llama_set_sampler(llama_context * ctx, llama_seq_id seq_id, llama_sampler * smpl) { return ctx->set_sampler(seq_id, smpl); } diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b562..c083e84eaffb 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -90,6 +90,16 @@ struct llama_context { float * get_embeddings_layer_inp(uint32_t lid); + // persistent device tensor holding the concat of the enabled layer-input tensors + // (zero-copy embd_dev path); null when not in use + ggml_tensor * get_embd_layer_inp_fused() { return embd_layer_inp_fused; } + + // event recorded after each compute that writes the fused tensor; null when disabled + ggml_backend_event_t get_embd_layer_inp_fused_event() { return embd_layer_inp_fused_event; } + + // persistent device tensor holding the encoder output (t_h_nextn); null when disabled + ggml_tensor * get_embd_nextn_persist() { return embd_nextn_persist; } + llama_token * get_sampled_tokens() const; llama_token get_sampled_token_ith(int32_t idx); @@ -138,7 +148,8 @@ struct llama_context { const llama_ubatch & ubatch, llm_graph_type gtype, llama_memory_context_i * mctx, - ggml_status & ret); + ggml_status & ret, + int64_t token_offset = 0); int encode(const llama_batch & batch_inp); int decode(const llama_batch & batch_inp); @@ -234,6 +245,15 @@ struct llama_context { // from backend into host-side embd_layer_inp buffers void extract_layer_inputs(const llm_graph_result * res, size_t token_offset, size_t n_tokens); + // lazily allocate the persistent device buffer that receives the concat of the + // enabled layer-input tensors (zero-copy embd_dev path). returns null when no + // layer-input extraction is enabled (host path). + ggml_tensor * ensure_embd_layer_inp_fused(); + + // lazily allocate the persistent device buffer that receives the encoder output + // (t_h_nextn) so the decoder KV-injection can alias it (zero-copy nextn path). + ggml_tensor * ensure_embd_nextn_persist(); + // // graph // @@ -258,7 +278,8 @@ struct llama_context { llm_graph_result * res, const llama_ubatch & ubatch, const llama_memory_context_i * mctx, - llm_graph_type gtype) const; + llm_graph_type gtype, + int64_t token_offset = 0) const; llm_graph_cb graph_get_cb() const; @@ -370,6 +391,23 @@ struct llama_context { // host buffer for the model output (logits and embeddings) ggml_backend_buffer_ptr buf_output; + // persistent device buffer for the concat of the enabled layer-input tensors + // (zero-copy embd_dev path). lazily allocated on first use; null when disabled. + ggml_tensor * embd_layer_inp_fused = nullptr; + ggml_context * embd_layer_inp_fused_ctx = nullptr; + ggml_backend_buffer_ptr embd_layer_inp_fused_buf; + + // recorded after each compute that writes embd_layer_inp_fused; lets a foreign + // context (the speculative draft) wait on the GPU stream instead of host-syncing + ggml_backend_event_t embd_layer_inp_fused_event = nullptr; + + // persistent device buffer holding the encoder output (t_h_nextn) of the previous + // compute call, so the DFlash decoder KV-injection can alias it via a view instead + // of a host read + H2D copy. lazily allocated on first use; null when disabled. + ggml_tensor * embd_nextn_persist = nullptr; + ggml_context * embd_nextn_persist_ctx = nullptr; + ggml_backend_buffer_ptr embd_nextn_persist_buf; + // keep copies of the per-sequence memory on the device std::map mem_storage; diff --git a/src/llama-ext.h b/src/llama-ext.h index 35d6e58adfa8..c38ffa3d044e 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -114,6 +114,26 @@ LLAMA_API void llama_set_embeddings_layer_inp(struct llama_context * ctx, uint32 // LLAMA_API float * llama_get_embeddings(struct llama_context * ctx); LLAMA_API float * llama_get_embeddings_layer_inp(struct llama_context * ctx, uint32_t lid); +// returns the persistent device tensor holding the concat of the enabled layer-input +// tensors (zero-copy embd_dev path). the tensor is [K*n_embd, n_batch] where K is the +// number of enabled embeddings_layer_inp layers. its contents are overwritten by the +// next llama_encode/llama_decode on the same context. returns null when zero-copy is +// not in use. +LLAMA_API const struct ggml_tensor * llama_get_embeddings_layer_inp_tensor(struct llama_context * ctx); + +// wait (GPU-side, no host block) until the target context's fused layer-input write +// from its last decode is visible to the draft context's backends. the consumer must +// call this before reading/consuming the tensor returned by +// llama_get_embeddings_layer_inp_tensor on the GPU. +LLAMA_API void llama_embd_layer_inp_wait(struct llama_context * ctx_tgt, struct llama_context * ctx_dft); + +// returns the persistent device tensor holding the encoder output (t_h_nextn) of the +// last encoder compute on this context (zero-copy nextn path). the tensor is +// [n_embd_out, n_batch]. its contents are overwritten by the next encode on the same +// context. returns null when the zero-copy nextn path is disabled. +LLAMA_API const struct ggml_tensor * llama_get_embeddings_nextn_tensor(struct llama_context * ctx); + + LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx); // diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 1896758c5da5..7edce7e45303 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -72,6 +72,12 @@ void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) { ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens*ggml_element_size(tokens)); } + if (ubatch->embd_dev) { + // zero-copy path: the input is a view of an external device tensor + // (already on the backend), nothing to upload + return; + } + if (ubatch->embd) { GGML_ASSERT(n_embd == embd->ne[0]); @@ -86,6 +92,9 @@ bool llm_graph_input_embd::can_reuse(const llm_graph_params & params) { res &= (!params.ubatch.token) || (tokens && tokens->ne[0] == params.ubatch.n_tokens); res &= (!params.ubatch.embd) || (embd && embd->ne[1] == params.ubatch.n_tokens); + // zero-copy path: the graph embeds a view of the external device tensor, so + // the tensor must be identical (pointer identity) for the graph to be reused + res &= (!params.ubatch.embd_dev) || (embd_dev_ptr == params.ubatch.embd_dev); return res; } @@ -93,6 +102,11 @@ bool llm_graph_input_embd::can_reuse(const llm_graph_params & params) { void llm_graph_input_embd_h::set_input(const llama_ubatch * ubatch) { const int64_t n_tokens = ubatch->n_tokens; + if (ubatch->embd_dev) { + // zero-copy path: inputs are views of external device tensors, nothing to upload + return; + } + if (ubatch->token) { ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens*ggml_element_size(tokens)); } else { @@ -119,6 +133,8 @@ bool llm_graph_input_embd_h::can_reuse(const llm_graph_params & params) { res &= (!params.ubatch.token) || (tokens && tokens->ne[0] == params.ubatch.n_tokens); res &= (!params.ubatch.embd) || (embd && embd->ne[1] == params.ubatch.n_tokens); res &= (!params.ubatch.embd) || (h && h->ne[1] == params.ubatch.n_tokens); + // zero-copy path: the graph embeds views of the external device tensor(s) + res &= (!params.ubatch.embd_dev) || (embd_dev_ptr == params.ubatch.embd_dev); return res; } @@ -2294,9 +2310,23 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const { ggml_set_input(inp->tokens); res->t_inp_tokens = inp->tokens; - inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, ubatch.n_tokens); - cb(inp->embd, "inp_embd", -1); - ggml_set_input(inp->embd); + if (ubatch.embd_dev) { + // zero-copy path: alias the external device tensor via a view (no allocation, + // no upload). the tensor is [n_embd, n_tokens] (contiguous) - view a column + // range starting at embd_dev_off. + GGML_ASSERT(n_embd_inp == ubatch.embd_dev->ne[0]); + GGML_ASSERT(ubatch.embd_dev_off + ubatch.n_tokens <= ubatch.embd_dev->ne[1]); + + inp->embd = ggml_view_2d(ctx0, ubatch.embd_dev, n_embd_inp, ubatch.n_tokens, + ubatch.embd_dev->nb[1], (size_t) ubatch.embd_dev_off * ubatch.embd_dev->nb[1]); + inp->embd_dev_ptr = ubatch.embd_dev; + // NOTE: no ggml_set_input - the tensor is already resident on the backend; + // flagging it as INPUT would make the scheduler copy it in. + } else { + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, ubatch.n_tokens); + cb(inp->embd, "inp_embd", -1); + ggml_set_input(inp->embd); + } // select one of the 2 inputs, based on the batch contents // ref: https://github.com/ggml-org/llama.cpp/pull/18550 diff --git a/src/llama-graph.h b/src/llama-graph.h index 94324c7457ed..8738e6f6ed2c 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -134,6 +134,10 @@ class llm_graph_input_embd : public llm_graph_input_i { ggml_tensor * tokens = nullptr; // I32 [n_batch] ggml_tensor * embd = nullptr; // F32 [n_embd, n_batch] + // external device tensor this graph aliases via a view (zero-copy path), if any + // used to invalidate the graph on pointer change + ggml_tensor * embd_dev_ptr = nullptr; + const int64_t n_embd = 0; }; @@ -151,6 +155,9 @@ class llm_graph_input_embd_h : public llm_graph_input_i { ggml_tensor * embd = nullptr; // F32 [n_embd, n_batch] ggml_tensor * h = nullptr; // F32 [n_embd, n_batch] + // external device tensor this graph aliases via a view (zero-copy path), if any + ggml_tensor * embd_dev_ptr = nullptr; + const int64_t n_embd = 0; }; @@ -753,6 +760,19 @@ struct llm_graph_params { const llama_memory_context_i * mctx; const llama_cross * cross; + // offset of this ubatch within the full batch (in tokens); used when writing the + // fused layer-input tensor so that multi-ubatch decodes land at the right columns. + int64_t token_offset = 0; + + // persistent device buffer that receives the concat of the enabled layer-input + // tensors (t_layer_inp[il]) each compute call; used by the zero-copy embd_dev path. + // null when zero-copy is not in use. + ggml_tensor * embd_layer_inp_fused = nullptr; + + // persistent device buffer that receives the encoder output (t_h_nextn) so the + // DFlash decoder KV-injection can alias it (zero-copy nextn path). null when disabled. + ggml_tensor * embd_nextn_persist = nullptr; + std::map samplers; static bool samplers_equal( @@ -790,7 +810,11 @@ struct llm_graph_params { (!ubatch.token && !other.ubatch.token) || (!ubatch.embd && !other.ubatch.embd) || (ubatch.token && other.ubatch.token && ubatch.embd && other.ubatch.embd) - ); + ) && + // zero-copy path: the graph embeds a view of the external device tensor, + // so both the tensor pointer and the column offset must match to reuse + ubatch.embd_dev == other.ubatch.embd_dev && + ubatch.embd_dev_off == other.ubatch.embd_dev_off; // when we split the batch using "equal_seqs" we have to verify that the participating sequences are the same // the reason is because the set of attention streams would be different for different sequences diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 5caa05e8b07d..842663e78d92 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -83,7 +83,7 @@ static llama_ubatch dsv4_build_raw_write_ubatch(const llama_ubatch & ubatch) { if (!dsv4_ubatch_has_coupled(ubatch)) { return ubatch; } - if (ubatch.embd) { + if (ubatch.embd || ubatch.embd_dev) { throw std::runtime_error("DSV4 coupled embedding ubatches are not supported"); } @@ -158,6 +158,8 @@ static llama_ubatch dsv4_build_raw_write_ubatch(const llama_ubatch & ubatch) { /*.n_pos =*/ ubatch.n_pos, /*.token =*/ data->token.empty() ? nullptr : data->token.data(), /*.embd =*/ nullptr, + /*.embd_dev =*/ nullptr, + /*.embd_dev_off =*/ 0, /*.pos =*/ data->pos.data(), /*.n_seq_id =*/ data->n_seq_id.data(), /*.seq_id =*/ data->seq_id.data(), diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 0d74a2135b6e..f4877515367d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2468,6 +2468,59 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const { llm->res->set_outputs(params); + // zero-copy embd_dev path: concat the enabled layer-input tensors into a + // persistent device buffer so a foreign context (e.g. the speculative draft) + // can alias it via a view on the next compute call. + if (params.embd_layer_inp_fused) { + const auto & enabled = params.cparams.embeddings_layer_inp; + + ggml_tensor * fused = nullptr; + for (size_t il = 0; il < enabled.size(); ++il) { + if (!enabled[il]) { + continue; + } + ggml_tensor * layer_inp = llm->res->get_layer_inp((int) il); + if (!layer_inp) { + GGML_ABORT("layer input tensor not found for embd_layer_inp_fused"); + } + fused = fused ? ggml_concat(llm->ctx0, fused, layer_inp, 0) : layer_inp; + } + + if (fused) { + GGML_ASSERT(fused->ne[0] == params.embd_layer_inp_fused->ne[0]); + // copy only the columns for this ubatch (n_tokens <= n_batch) into a view + // of the persistent buffer at the correct offset; the rest is untouched. + ggml_tensor * fused_dst = ggml_view_2d( + llm->ctx0, params.embd_layer_inp_fused, + params.embd_layer_inp_fused->ne[0], params.ubatch.n_tokens, + params.embd_layer_inp_fused->nb[1], + (size_t) params.token_offset * params.embd_layer_inp_fused->nb[1]); + ggml_tensor * dst = ggml_cpy(llm->ctx0, fused, fused_dst); + ggml_build_forward_expand(llm->res->get_gf(), dst); + } + } + + // zero-copy nextn path: copy the encoder output (t_h_nextn) into a persistent + // device buffer so the DFlash decoder KV-injection can alias it via a view, + // avoiding the host read + H2D copy in the speculative loop. + if (params.gtype == LLM_GRAPH_TYPE_ENCODER && params.embd_nextn_persist) { + ggml_tensor * h_nextn = llm->res->get_h_nextn(); + if (h_nextn) { + // the encoder output is always consumed by the decoder KV-injection with + // embd_dev_off = 0, so the persist write must land at column 0. encode runs + // as a single ubatch with token_offset == 0; assert that invariant. + GGML_ASSERT(params.token_offset == 0 && "encoder nextn persist requires token_offset == 0"); + GGML_ASSERT(h_nextn->ne[0] == params.embd_nextn_persist->ne[0]); + ggml_tensor * nextn_dst = ggml_view_2d( + llm->ctx0, params.embd_nextn_persist, + params.embd_nextn_persist->ne[0], params.ubatch.n_tokens, + params.embd_nextn_persist->nb[1], + (size_t) params.token_offset * params.embd_nextn_persist->nb[1]); + ggml_tensor * dst = ggml_cpy(llm->ctx0, h_nextn, nextn_dst); + ggml_build_forward_expand(llm->res->get_gf(), dst); + } + } + return llm->res->get_gf(); } diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 5b70a5179496..785613e376fc 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -209,8 +209,21 @@ template <> ggml_tensor * llama_model_dflash::graph::build_inp_embd_enc() const { auto inp_target = std::make_unique(hparams.n_embd_inp_enc()); - inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp_enc(), n_tokens); - ggml_set_input(inp_target->embd); + if (ubatch.embd_dev) { + // zero-copy path: alias the external fused device tensor (target features + // concat of the extract layers) via a view - no allocation, no upload. + GGML_ASSERT(hparams.n_embd_inp_enc() == ubatch.embd_dev->ne[0]); + GGML_ASSERT(ubatch.embd_dev_off + ubatch.n_tokens <= ubatch.embd_dev->ne[1]); + + inp_target->embd = ggml_view_2d(ctx0, ubatch.embd_dev, hparams.n_embd_inp_enc(), ubatch.n_tokens, + ubatch.embd_dev->nb[1], (size_t) ubatch.embd_dev_off * ubatch.embd_dev->nb[1]); + inp_target->embd_dev_ptr = ubatch.embd_dev; + // NOTE: no ggml_set_input - the tensor is already resident on the backend; + // flagging it as INPUT would make the scheduler copy it in. + } else { + inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp_enc(), n_tokens); + ggml_set_input(inp_target->embd); + } ggml_tensor * cur = inp_target->embd; cb(cur, "inp_embd", -1); @@ -369,12 +382,23 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); - // KV cache injection - if (ubatch.embd) { + // KV cache injection (host embeddings or zero-copy device alias) + if (ubatch.embd || ubatch.embd_dev) { auto inp = std::make_unique(n_embd); - inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens); - ggml_set_input(inp->embd); + if (ubatch.embd_dev) { + // zero-copy path: alias the persistent encoder-output tensor via a view + GGML_ASSERT(n_embd == ubatch.embd_dev->ne[0]); + GGML_ASSERT(ubatch.embd_dev_off + ubatch.n_tokens <= ubatch.embd_dev->ne[1]); + + inp->embd = ggml_view_2d(ctx0, ubatch.embd_dev, n_embd, ubatch.n_tokens, + ubatch.embd_dev->nb[1], (size_t) ubatch.embd_dev_off * ubatch.embd_dev->nb[1]); + inp->embd_dev_ptr = ubatch.embd_dev; + // NOTE: no ggml_set_input - the tensor is already resident on the backend + } else { + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(inp->embd); + } ggml_tensor * inp_g = inp->embd; cb(inp_g, "inp_g_embeddings", -1); @@ -571,7 +595,7 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_ llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa(); // KV cache injection: fused target features from the encoder - if (ubatch.embd) { + if (ubatch.embd || ubatch.embd_dev) { auto inp = std::make_unique(n_embd); inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens); diff --git a/src/models/eagle3.cpp b/src/models/eagle3.cpp index be466056df69..5ec6845b6030 100644 --- a/src/models/eagle3.cpp +++ b/src/models/eagle3.cpp @@ -119,8 +119,19 @@ ggml_tensor * llama_model_eagle3::graph::build_inp_embd_enc() const { // Input: Target model features (3 layers concatenated: low, mid, high) // Data will be provided via ubatch->embd in encode_eagle3_features() auto inp_target = std::make_unique(hparams.n_embd_inp_enc()); - inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp_enc(), n_tokens); - ggml_set_input(inp_target->embd); + if (ubatch.embd_dev) { + // zero-copy path: alias the external fused device tensor via a view + GGML_ASSERT(hparams.n_embd_inp_enc() == ubatch.embd_dev->ne[0]); + GGML_ASSERT(ubatch.embd_dev_off + ubatch.n_tokens <= ubatch.embd_dev->ne[1]); + + inp_target->embd = ggml_view_2d(ctx0, ubatch.embd_dev, hparams.n_embd_inp_enc(), ubatch.n_tokens, + ubatch.embd_dev->nb[1], (size_t) ubatch.embd_dev_off * ubatch.embd_dev->nb[1]); + inp_target->embd_dev_ptr = ubatch.embd_dev; + // NOTE: no ggml_set_input - the tensor is already resident on the backend + } else { + inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp_enc(), n_tokens); + ggml_set_input(inp_target->embd); + } cur = inp_target->embd; cb(cur, "inp_embd", -1); diff --git a/src/models/glm4-moe.cpp b/src/models/glm4-moe.cpp index d60e47ddf0c6..53d4f6564027 100644 --- a/src/models/glm4-moe.cpp +++ b/src/models/glm4-moe.cpp @@ -144,7 +144,7 @@ llama_model_glm4_moe::graph::graph(const llama_model & model, const llm_graph_pa inpL = build_inp_embd(model.tok_embd); bool use_mrope = hparams.use_mrope(); - if (ubatch.embd && !use_mrope) { + if ((ubatch.embd || ubatch.embd_dev) && !use_mrope) { // unfortunately, we need to forcefully stop here, to avoid users complaining about wrong results GGML_ABORT("This GGUF does not support multimodal. Please reconvert it."); } diff --git a/src/models/glm4.cpp b/src/models/glm4.cpp index b4326c5f2107..289c3e07cf89 100644 --- a/src/models/glm4.cpp +++ b/src/models/glm4.cpp @@ -83,7 +83,7 @@ llama_model_glm4::graph::graph(const llama_model & model, const llm_graph_params inpL = build_inp_embd(model.tok_embd); bool use_mrope = hparams.use_mrope(); - if (ubatch.embd && !use_mrope) { + if ((ubatch.embd || ubatch.embd_dev) && !use_mrope) { // unfortunately, we need to forcefully stop here, to avoid users complaining about wrong results GGML_ABORT("This GGUF does not support multimodal. Please reconvert it."); } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 842e4203cd2a..582f79cc79d3 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -303,6 +303,11 @@ struct server_slot { common_sampler_ptr smpl; + // persistent sampler snapshot used to roll back the sampler state on partial + // draft acceptance. initialized once alongside smpl; we copy state into it + // (llama_sampler_copy, no allocation) instead of deep-cloning every token. + common_sampler_ptr smpl_save; + llama_token sampled; // in speculative mode, this is the last accepted token // for TTS models, this is the embd generated from prev step, decode this to generate next hidden state @@ -1721,6 +1726,8 @@ struct server_context_impl { if (task.need_sampling()) { try { slot.smpl.reset(common_sampler_init(model_tgt, task.params.sampling)); + // one persistent snapshot for speculative rollback (state copied, not cloned) + slot.smpl_save.reset(common_sampler_init(model_tgt, task.params.sampling)); } catch (std::exception & e) { std::string err_msg = std::string("Failed to initialize samplers: ") + e.what(); send_error(task, err_msg, ERROR_TYPE_INVALID_REQUEST); @@ -3814,7 +3821,11 @@ struct server_context_impl { // verify and try to accept the draft { - common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get())); + // snapshot the sampler state (cheap copy_state, no deep clone) so we + // can restore it on partial draft acceptance + if (slot.smpl_save) { + common_sampler_copy(slot.smpl.get(), slot.smpl_save.get()); + } GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); @@ -3852,7 +3863,9 @@ struct server_context_impl { slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1); slot.prompt.tokens.keep_first(ckpt.n_tokens); - common_sampler_copy(smpl_save.get(), slot.smpl.get()); + if (slot.smpl_save) { + common_sampler_copy(slot.smpl_save.get(), slot.smpl.get()); + } return; }