vrr/koboldcpp
2
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-08-21 14:25:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

spec: enable backend sampling for both dflash & dspark (#26958)

Browse source 
* dflash: enable backend sampling for both dflash & dspark

* enable p_min > 0 in backend sampling and add guard

* cont : add TODO

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
Ruixiang Wang 2026-08-13 11:07:28 +02:00 • committed by GitHub
parent eeae28b67e
commit 0d0bfcd4fd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 48 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

48
common/speculative.cpp
View file

@ -912,6 +912,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
std::vector<common_sampler_ptr> smpls;
// backend sampler chain per seq, attached to ctx_dft
std::vector<llama_sampler *> 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
@ -985,6 +988,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(10));
if (!llama_set_sampler(ctx_dft, seq_id, chain)) {
SPC_WRN("backend offload failed for seq_id=%d; using CPU sampler\n", (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);
@ -995,6 +1014,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);
}
@ -2263,6 +2294,23 @@ common_params common_base_params_to_speculative(const common_params & params) {
result.n_outputs_max = params.n_parallel;
result.n_outputs_max_per_seq = 1;
// dflash/dspark decode the whole noise block in a single pass and sample every block position on the backend
// TODO: refactor such properties to be announced by the speculative types
// something like `struct common_speculative_type_props common_speculative_type_get_props(...);`
const bool has_block_draft = std::any_of(
params.speculative.types.begin(), params.speculative.types.end(),
[](common_speculative_type t) {
return t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;
});
if (has_block_draft) {
// per-seq output positions: DFlash decodes anchor + n_max masks (n_max + 1); DSpark n_max -> +1 covers both
const int32_t per_seq = std::max(1, params_spec.n_max + 1);
result.n_outputs_max = params.n_parallel * per_seq;
if (params_spec.backend_sampling) {
result.n_outputs_max_per_seq = per_seq;
}
}
return result;
}