diff --git a/common/speculative.cpp b/common/speculative.cpp index 3951bbed5455..3a445f8da1e1 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -908,6 +908,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { std::vector smpls; + // backend sampler chain per seq, attached to ctx_dft + std::vector backend_chains; + int32_t n_embd_dec = 0; // draft hidden size int32_t n_embd_enc = 0; // target_layer_ids_n * target_hidden_size int32_t n_embd_tgt = 0; // target model hidden size @@ -951,7 +954,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); LOG_INF("%s: adding speculative implementation 'draft-dflash'\n", __func__); - LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min); + LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f, backend_sampling=%d\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min, (int) this->params.backend_sampling); LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u\n", __func__, block_size, mask_token_id, target_layer_ids_n); // DFlash input is [id_last, * (block_size-1)], so it can draft at most block_size-1 tokens per step @@ -973,6 +976,22 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { s.reset(common_sampler_init(model_dft, sparams)); } + // offload draft sampling to the backend + backend_chains.assign(n_seq, nullptr); + if (this->params.backend_sampling) { + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { + llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params()); + llama_sampler_chain_add(chain, llama_sampler_init_top_k(1)); + + if (!llama_set_sampler(ctx_dft, seq_id, chain)) { + LOG_WRN("%s: backend offload failed for seq_id=%d; using CPU sampler\n", __func__, (int) seq_id); + llama_sampler_free(chain); + chain = nullptr; + } + backend_chains[seq_id] = chain; + } + } + // turn on extraction of the target layers' input embeddings for (uint32_t k = 0; k < target_layer_ids_n; ++k) { llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true); @@ -983,6 +1002,18 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } ~common_speculative_impl_draft_dflash() override { + auto * ctx_dft = this->params.ctx_dft; + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) backend_chains.size(); ++seq_id) { + if (backend_chains[seq_id] == nullptr) { + continue; + } + if (ctx_dft) { + llama_set_sampler(ctx_dft, seq_id, nullptr); + } + llama_sampler_free(backend_chains[seq_id]); + } + backend_chains.clear(); + llama_batch_free(batch); llama_batch_free(batch_inject); } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0465430df43a..bb8a2a5d09bc 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1565,9 +1565,8 @@ static std::map build_seq_to_output_row(const llama_ubat continue; } - const llama_seq_id seq_id = ubatch.seq_id[i][0]; // row_offset is the number of output tokens before this ubatch. - seq_to_row[seq_id] = row_offset + local; + seq_to_row[local] = row_offset + local; ++local; } return seq_to_row; @@ -1704,29 +1703,6 @@ int llama_context::decode(const llama_batch & batch_inp) { const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max; - // TODO: avoid this workaround in the future - if (has_samplers && batch_inp.logits) { - std::vector seq_output_count(n_seq_max, 0); - - for (int32_t i = 0; i < batch_inp.n_tokens; ++i) { - if (batch_inp.logits[i] == 0) { - continue; - } - - const int ns = batch_inp.n_seq_id ? batch_inp.n_seq_id[i] : 1; - - for (int32_t s = 0; s < ns; ++s) { - const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0; - - seq_output_count[seq_id]++; - if (seq_output_count[seq_id] > 1) { - LLAMA_LOG_ERROR("%s: backend sampling requires at most one output token per sequence (seq_id %d had %d)\n", - __func__, seq_id, seq_output_count[seq_id]); - return -1; - } - } - } - } if (!balloc->init(batch_inp, vocab, memory.get(), n_embd, n_seq_max, output_all)) { LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 4c86e43c1f74..c07b0aa9c7f2 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3398,13 +3398,14 @@ void llm_graph_context::build_sampling() const { auto inp_sampling = std::make_unique(samplers); res->add_input(std::move(inp_sampling)); - std::map seq_to_logit_row; + // a sequence may have more than one output row per decode (e.g. DFlash blocks) + std::map> seq_to_logit_rows; int32_t logit_row_idx = 0; for (uint32_t i = 0; i < ubatch.n_tokens; i++) { if (ubatch.output[i]) { llama_seq_id seq_id = ubatch.seq_id[i][0]; - seq_to_logit_row[seq_id] = logit_row_idx; + seq_to_logit_rows[seq_id].push_back(logit_row_idx); logit_row_idx++; } } @@ -3418,47 +3419,48 @@ void llm_graph_context::build_sampling() const { ggml_tensor * logits_t = ggml_pad(ctx0, res->t_logits, 0, 1, 0, 0); for (const auto & [seq_id, sampler] : samplers) { - const auto it = seq_to_logit_row.find(seq_id); + const auto it = seq_to_logit_rows.find(seq_id); - // inactive samplers always work on the first row - const auto row_idx = it != seq_to_logit_row.end() ? it->second : 0; - const int i_out = it != seq_to_logit_row.end() ? 1 : 0; + const std::vector rows = it != seq_to_logit_rows.end() ? it->second : std::vector{ 0 }; + const int i_out = it != seq_to_logit_rows.end() ? 1 : 0; - ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], row_idx * logits_t->nb[1]); - ggml_format_name(logits_seq, "logits_seq_%d", seq_id); + for (const int32_t row_idx : rows) { + ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], row_idx * logits_t->nb[1]); + ggml_format_name(logits_seq, "logits_seq_%d", seq_id); - struct llama_sampler_data data = { - /*.logits =*/ logits_seq, - /*.probs =*/ nullptr, - /*.sampled =*/ nullptr, - /*.candidates =*/ nullptr, - }; + struct llama_sampler_data data = { + /*.logits =*/ logits_seq, + /*.probs =*/ nullptr, + /*.sampled =*/ nullptr, + /*.candidates =*/ nullptr, + }; - assert(sampler->iface->backend_apply); - sampler->iface->backend_apply(sampler, ctx0, gf, &data); + assert(sampler->iface->backend_apply); + sampler->iface->backend_apply(sampler, ctx0, gf, &data); - if (data.sampled != nullptr) { - res->t_sampled[seq_id] = data.sampled; - outs[1] = data.sampled; - ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); - } + if (data.sampled != nullptr) { + if (i_out) res->t_sampled[row_idx] = data.sampled; + outs[1] = data.sampled; + ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); + } - if (data.probs != nullptr) { - res->t_sampled_probs[seq_id] = data.probs; - outs[1] = data.probs; - ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); - } + if (data.probs != nullptr) { + if (i_out) res->t_sampled_probs[row_idx] = data.probs; + outs[1] = data.probs; + ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); + } - if (data.logits != nullptr) { - res->t_sampled_logits[seq_id] = data.logits; - outs[1] = data.logits; - ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); - } + if (data.logits != nullptr) { + if (i_out) res->t_sampled_logits[row_idx] = data.logits; + outs[1] = data.logits; + ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); + } - if (data.candidates != nullptr) { - res->t_candidates[seq_id] = data.candidates; - outs[1] = data.candidates; - ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); + if (data.candidates != nullptr) { + if (i_out) res->t_candidates[row_idx] = data.candidates; + outs[1] = data.candidates; + ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); + } } }