Skip to content

dflash: port #219 onto feature (rebased full CL) #220

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
TheTom wants to merge 1 commit into
base: feature/turboquant-kv-cache
Choose a base branch
from
Open

dflash: port #219 onto feature (rebased full CL) #220

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions common/arg.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3600,6 +3600,27 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
}
}
).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE"));
add_opt(common_arg(
{"--eagle3"},
"use EAGLE3 speculative decoding with the draft model",
[](common_params & params) {
params.speculative.draft.eagle3 = true;
}
).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
add_opt(common_arg(
{"--dflash"},
"use DFlash speculative decoding with the draft model",
[](common_params & params) {
params.speculative.draft.dflash = true;
}
).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
add_opt(common_arg(
{"-cd", "--ctx-size-draft"}, "N",
string_format("size of the prompt context for the draft model (default: %d, 0 = loaded from model)", params.speculative.draft.n_ctx),
[](common_params & params, int value) {
params.speculative.draft.n_ctx = value;
}
).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CTX"));

add_opt(common_arg(
{"--spec-draft-n-max"}, "N",
Expand Down
5 changes: 5 additions & 0 deletions common/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -325,6 +325,11 @@ struct common_params_speculative_draft {
std::vector<ggml_backend_dev_t> devices; // devices to use for offloading

std::vector<llama_model_tensor_buft_override> tensor_buft_overrides;

bool eagle3 = false; // use EAGLE3 speculative decoding
bool dflash = false; // use DFlash speculative decoding
int32_t n_ctx = 0; // draft context size

};

struct common_params_speculative_ngram_mod {
Expand Down
106 changes: 88 additions & 18 deletions common/speculative.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -835,6 +835,14 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

// scratch buffer for concatenated target features [n_tokens, n_embd_enc]
std::vector<float> features_buf;
// Stashed encoder outputs for deferred KV cache injection
struct StashedG {
llama_pos pos;
std::vector<float> data; // n_embd_dec floats
};
std::vector<std::vector<StashedG>> stashed; // [n_seq][...]
static constexpr int32_t MAX_STASH = 64;
bool m_use_deferred = true;

common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq)
: common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, n_seq)
Expand Down Expand Up @@ -869,7 +877,17 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
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: - 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, <mask> * (block_size-1)], so it can draft at most block_size-1 tokens per step
// Detect target model type: Gemma4 models regress with deferred injection (low acceptance rate)
{
char model_desc[128] = {};
llama_model_desc(model_tgt, model_desc, sizeof(model_desc));
if (strstr(model_desc, "gemma")) {
m_use_deferred = false;
LOG_INF("%s: - deferred_kv_injection=0 (disabled for %s)\n", __func__, model_desc);
} else {
LOG_INF("%s: - deferred_kv_injection=1 (enabled for %s)\n", __func__, model_desc);
}
}
if (this->params.n_max > block_size - 1 || this->params.n_min > block_size - 1) {
LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained DFlash block size %d -- clamping to %d\n",
__func__, this->params.n_max, this->params.n_min, block_size, block_size - 1);
Expand All @@ -880,6 +898,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq);
batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq);

stashed.resize(n_seq);

smpls.resize(n_seq);
for (auto & s : smpls) {
common_params_sampling sparams;
Expand All @@ -889,13 +909,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
s.reset(common_sampler_init(model_dft, sparams));
}

// turn on extraction of the target layers' input embeddings
// Enable extraction of target model layer outputs for DFlash encoder
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);
}

// Enable nextn (encoder) output on the draft context for DFlash decoder K/V injection
llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention
}

~common_speculative_impl_draft_dflash() override {
Expand Down Expand Up @@ -996,31 +1016,74 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
const float * inp_g = llama_get_embeddings_nextn(ctx_dft);
GGML_ASSERT(inp_g && "DFlash encoder produced no output.");

// 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));

for (int32_t i = 0; i < n_chunk; ++i) {
batch_inject.pos[i] = batch_in.pos[i_batch_beg[seq_id] + offset + i];
batch_inject.n_seq_id[i] = 1;
batch_inject.seq_id[i][0] = seq_id;
batch_inject.logits[i] = false;
}
rc = llama_decode(ctx_dft, batch_inject);
if (rc != 0) {
LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
__func__, rc, (int) n_chunk, (int) offset);
return false;
if (m_use_deferred) {
// stash encoder output for deferred KV cache injection
auto & stash = stashed[seq_id];
for (int32_t i = 0; i < n_chunk; ++i) {
StashedG sg;
sg.pos = batch_in.pos[i_batch_beg[seq_id] + offset + i];
sg.data.assign(inp_g + (size_t) i * n_embd_dec, inp_g + (size_t) (i + 1) * n_embd_dec);
stash.push_back(std::move(sg));
}
// auto-flush to prevent unbounded accumulation
if ((int32_t) stash.size() >= MAX_STASH) {
flush_injection(seq_id);
}
} else {
// immediate injection (original per-chunk behavior)
auto & stash = stashed[seq_id];
for (int32_t i = 0; i < n_chunk; ++i) {
StashedG sg;
sg.pos = batch_in.pos[i_batch_beg[seq_id] + offset + i];
sg.data.assign(inp_g + (size_t) i * n_embd_dec, inp_g + (size_t) (i + 1) * n_embd_dec);
stash.push_back(std::move(sg));
}
flush_injection(seq_id);
}
}
}

return true;
}

// Batch-inject all stashed encoder outputs into the draft decoder's KV cache
void flush_injection(llama_seq_id seq_id) {
auto & stash = stashed[seq_id];
if (stash.empty()) return;

const int32_t n = (int32_t) stash.size();
batch_inject.n_tokens = n;
for (int32_t i = 0; i < n; ++i) {
std::memcpy(batch_inject.embd + (size_t) i * n_embd_dec,
stash[i].data.data(), (size_t) n_embd_dec * sizeof(float));
batch_inject.pos[i] = stash[i].pos;
batch_inject.n_seq_id[i] = 1;
batch_inject.seq_id[i][0] = seq_id;
batch_inject.logits[i] = false;
}

auto * ctx_dft = params.ctx_dft;
int32_t rc = llama_decode(ctx_dft, batch_inject);
if (rc != 0) {
LOG_ERR("%s: flush_injection llama_decode(ctx_dft) failed rc=%d (n_tokens=%d)\n",
__func__, rc, n);
}

stash.clear();
}

void draft(common_speculative_draft_params_vec & dparams) override {
auto & ctx_dft = params.ctx_dft;

if (m_use_deferred) {
// flush all stashed KV injections before drafting
for (llama_seq_id seq_id = 0; seq_id < n_seq; ++seq_id) {
if (dparams[seq_id].drafting) {
flush_injection(seq_id);
}
}
}

common_batch_clear(batch);

// build one batch holding every drafting sequence's noise block into a single decode)
Expand Down Expand Up @@ -2082,6 +2145,13 @@ common_speculative * common_speculative_init(common_params_speculative & params,
bool has_mtp = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_MTP)) && params.draft.ctx_dft != nullptr;
bool has_draft_dflash = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)) && params.draft.ctx_dft != nullptr;

// If --dflash or --eagle3 flags are set, enable the corresponding type
if (!has_draft_dflash && params.draft.dflash && params.draft.ctx_dft != nullptr) {
has_draft_dflash = true;
}
if (!has_draft_eagle3 && params.draft.eagle3 && params.draft.ctx_dft != nullptr) {
has_draft_eagle3 = true;
}


bool has_ngram_cache = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_NGRAM_CACHE));
Expand Down
2 changes: 2 additions & 0 deletions examples/speculative-simple/speculative-simple.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,8 @@ int main(int argc, char ** argv) {
}

auto cparams = common_context_params_to_llama(params_dft);
// DFlash/EAGLE3 draft models may reuse the target model's embeddings/output
cparams.ctx_other = ctx_tgt;
ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams));

params.speculative.draft.ctx_tgt = ctx_tgt;
Expand Down
Loading
Loading