* spec : add DFlash2 support (local convolution + candidate selector) (#27342) * support DFlash2 * Add p_min in DFlash2 Assisted-by: Claude Opus 5 * Revert unnecessary changes Assisted-by: Claude Opus 5 * Revert draft sampling in rejection sampling Assisted-by: Claude Opus 5 * Refactor code structure Assisted-by: Claude Opus 5 * Delete embedding scaling Assisted-by: Claude Opus 5 * Gate output transforms on DFlash2 Assisted-by: Claude Opus 5 * Optimize Dflash 2 cost Assisted-by: Claude Opus 5 * Avoid using atoi Assisted-by: Claude Opus 5 * Modify comments Assisted-by: Claude Opus 5 * Move llama_model_dflash_selector_top_k to llama-ext.h Assisted-by: Claude Opus 5 * Formatting Assisted-by: Claude Opus 5 * Apply patch to fix the mrope bug Assisted-by: Claude Opus 5 * fix ci Assisted-by: Claude Opus 5 * Fix graph number calculation Assisted-by: Claude Opus 5 * rename hid and unary Assisted-by: Claude Opus 5 --------- Co-authored-by: Jian Chen <jianchen0311@gmail.com> Co-authored-by: Xuan-Son Nguyen <son@huggingface.co> * revert top-k.cu changes --------- Co-authored-by: Zihan Zhang <tiancaizhangdaxian@sjtu.edu.cn> Co-authored-by: Jian Chen <jianchen0311@gmail.com>
This commit is contained in:
Xuan-Son Nguyen
co-authored by
Jian Chen
Zihan Zhang
Jian Chen
Zihan Zhang
parent
58546250cf
commit
b10f9ca58c
+69
-5
@@ -925,6 +925,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
int32_t block_size = 0;
|
||||
llama_token mask_token_id = 0;
|
||||
|
||||
bool is_dflash2 = false;
|
||||
bool is_mrope = false;
|
||||
int32_t selector_top_k = 0;
|
||||
|
||||
// draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
|
||||
const bool is_dspark;
|
||||
|
||||
@@ -969,6 +973,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
sample_from_anchor = std::strcmp(buf, "true") == 0;
|
||||
}
|
||||
}
|
||||
|
||||
selector_top_k = llama_model_dflash_selector_top_k(model_dft);
|
||||
is_dflash2 = selector_top_k > 0;
|
||||
mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft));
|
||||
|
||||
LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str());
|
||||
@@ -990,6 +997,13 @@ 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);
|
||||
|
||||
// embd batches on an M-RoPE draft need 4 position rows per token
|
||||
is_mrope = llama_model_rope_type(model_dft) == LLAMA_ROPE_TYPE_MROPE;
|
||||
if (is_mrope) {
|
||||
free(batch_inject.pos);
|
||||
batch_inject.pos = (llama_pos *) malloc(sizeof(llama_pos) * 4 * llama_n_batch(ctx_dft));
|
||||
}
|
||||
|
||||
smpls.resize(n_seq);
|
||||
for (auto & s : smpls) {
|
||||
common_params_sampling sparams;
|
||||
@@ -1001,7 +1015,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
|
||||
// offload draft sampling to the backend
|
||||
backend_chains.assign(n_seq, nullptr);
|
||||
if (this->params.backend_sampling) {
|
||||
if (this->params.backend_sampling && !is_dflash2) {
|
||||
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));
|
||||
@@ -1020,7 +1034,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
|
||||
}
|
||||
|
||||
llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
|
||||
// DFlash2 reads its selector lattice from h_nextn and never consumes raw logits.
|
||||
llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2);
|
||||
llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention
|
||||
}
|
||||
|
||||
@@ -1121,11 +1136,24 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
}
|
||||
|
||||
// fuse extracted features through DFlash encoder
|
||||
// M-RoPE drafts read 4 position rows per token from embd batches, so pass them explicitly
|
||||
std::vector<llama_pos> enc_pos;
|
||||
if (is_mrope) {
|
||||
enc_pos.resize((size_t) 4 * n_chunk);
|
||||
for (int32_t i = 0; i < n_chunk; ++i) {
|
||||
const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i];
|
||||
enc_pos[0 * n_chunk + i] = p;
|
||||
enc_pos[1 * n_chunk + i] = p;
|
||||
enc_pos[2 * n_chunk + i] = p;
|
||||
enc_pos[3 * n_chunk + i] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
llama_batch enc_batch = {
|
||||
/*.n_tokens =*/ n_chunk,
|
||||
/*.token =*/ nullptr,
|
||||
/*.embd =*/ features_buf.data(),
|
||||
/*.pos =*/ nullptr,
|
||||
/*.pos =*/ is_mrope ? enc_pos.data() : nullptr,
|
||||
/*.n_seq_id =*/ nullptr,
|
||||
/*.seq_id =*/ nullptr,
|
||||
/*.logits =*/ nullptr,
|
||||
@@ -1146,7 +1174,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
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];
|
||||
const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i];
|
||||
batch_inject.pos[i] = p;
|
||||
if (is_mrope) {
|
||||
batch_inject.pos[1 * n_chunk + i] = p;
|
||||
batch_inject.pos[2 * n_chunk + i] = p;
|
||||
batch_inject.pos[3 * n_chunk + i] = 0;
|
||||
}
|
||||
batch_inject.n_seq_id[i] = 1;
|
||||
batch_inject.seq_id[i][0] = seq_id;
|
||||
batch_inject.logits[i] = false;
|
||||
@@ -1189,7 +1223,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
i_block_beg[seq_id] = batch.n_tokens;
|
||||
n_block [seq_id] = n_block_tokens;
|
||||
for (int32_t i = 0; i < n_block_tokens; ++i) {
|
||||
common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true);
|
||||
common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, !is_dflash2);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1217,6 +1251,36 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
|
||||
auto & result = *dp.result;
|
||||
|
||||
if (is_dflash2) {
|
||||
const float * lattice = llama_get_embeddings_nextn(ctx_dft);
|
||||
GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");
|
||||
|
||||
int32_t predecessor = 0;
|
||||
for (int32_t i = 1; i < n_block_tokens; ++i) {
|
||||
const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
|
||||
const float * scores = row + selector_top_k + (size_t) predecessor * selector_top_k;
|
||||
|
||||
predecessor = (int32_t) std::distance(scores,
|
||||
std::max_element(scores, scores + selector_top_k));
|
||||
if (params.p_min > 0.0f) {
|
||||
// softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max))
|
||||
float sum = 0.0f;
|
||||
for (int32_t k = 0; k < selector_top_k; ++k) {
|
||||
sum += std::exp(scores[k] - scores[predecessor]);
|
||||
}
|
||||
if (1.0f / sum < params.p_min) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
result.push_back((llama_token) row[predecessor]);
|
||||
}
|
||||
|
||||
if (result.size() < (size_t) params.n_min) {
|
||||
result.clear();
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (is_dspark) {
|
||||
// DSpark: read from the first draft slot, truncate below the confidence threshold
|
||||
const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr;
|
||||
|
||||
Reference in New Issue
Block a user