Skip to content

Commit 28ad034

Browse files
ruixiang63ruixiang63
committed
dflash: first working POC
1 parent 5bb2d50 commit 28ad034

24 files changed

Lines changed: 829 additions & 11 deletions

‎common/arg.cpp‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3420,6 +3420,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
34203420
params.speculative.eagle3 = true;
34213421
}
34223422
).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_CLI}));
3423+
add_opt(common_arg(
3424+
{"--dflash"},
3425+
"use DFlash speculative decoding with the draft model",
3426+
[](common_params & params) {
3427+
params.speculative.dflash = true;
3428+
}
3429+
).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_CLI}));
34233430
add_opt(common_arg(
34243431
{"-cd", "--ctx-size-draft"}, "N",
34253432
string_format("size of the prompt context for the draft model (default: %d, 0 = loaded from model)", params.speculative.n_ctx),

‎common/common.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,7 @@ enum common_speculative_type {
170170
COMMON_SPECULATIVE_TYPE_NONE, // no speculative decoding
171171
COMMON_SPECULATIVE_TYPE_DRAFT, // draft model
172172
COMMON_SPECULATIVE_TYPE_EAGLE3, // eagle draft model
173+
COMMON_SPECULATIVE_TYPE_DFLASH, // dflash draft model
173174
COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding
174175
COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, // self-speculative decoding with n-gram keys only
175176
COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values
@@ -296,6 +297,7 @@ struct common_params_speculative {
296297
llama_context_params cparams_dft; // these are the parameters for the draft llama_context
297298

298299
bool eagle3 = false; // use EAGLE3 speculative decoding
300+
bool dflash = false; // use DFlash speculative decoding
299301

300302
int32_t n_ctx = 0; // draft context size
301303
int32_t n_gpu_layers = -1; // number of layers to store in VRAM for the draft model (-1 - use default)

‎common/speculative.cpp‎

Lines changed: 150 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ const std::vector<enum common_speculative_type> common_speculative_types = {
2121
COMMON_SPECULATIVE_TYPE_NONE,
2222
COMMON_SPECULATIVE_TYPE_DRAFT,
2323
COMMON_SPECULATIVE_TYPE_EAGLE3,
24+
COMMON_SPECULATIVE_TYPE_DFLASH,
2425
COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE,
2526
COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K,
2627
COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V,
@@ -32,6 +33,7 @@ const std::map<std::string, enum common_speculative_type> common_speculative_typ
3233
{"none", COMMON_SPECULATIVE_TYPE_NONE},
3334
{"draft", COMMON_SPECULATIVE_TYPE_DRAFT},
3435
{"eagle3", COMMON_SPECULATIVE_TYPE_EAGLE3},
36+
{"dflash", COMMON_SPECULATIVE_TYPE_DFLASH},
3537
{"ngram_simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE},
3638
{"ngram_map_k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K},
3739
{"ngram_map_k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V},
@@ -595,6 +597,139 @@ struct common_speculative_state_eagle3 : public common_speculative_state {
595597
}
596598
};
597599

600+
struct common_speculative_state_dflash : public common_speculative_state {
601+
llama_context * ctx_tgt;
602+
603+
common_sampler * smpl;
604+
605+
llama_batch batch;
606+
607+
struct llama_context * ctx_dft_enc = nullptr;
608+
struct llama_context * ctx_dft_dec = nullptr;
609+
610+
int32_t dflash_n_past = 0;
611+
612+
// Host-side buffer: accumulated DFlash-encoded target features across all
613+
// committed prompt+drafted tokens. Grows by `n_new * n_embd` floats per draft step
614+
// and is fed to the DFlash decoder via llama_set_dflash_accumulated_target_ctx()
615+
std::vector<float> accumulated_ctx;
616+
617+
common_speculative_state_dflash(
618+
enum common_speculative_type type,
619+
llama_context * ctx_tgt,
620+
llama_context * ctx_dft_enc,
621+
llama_context * ctx_dft_dec)
622+
: common_speculative_state(type)
623+
, ctx_tgt(ctx_tgt)
624+
, ctx_dft_enc(ctx_dft_enc)
625+
, ctx_dft_dec(ctx_dft_dec)
626+
{
627+
batch = llama_batch_init(llama_n_batch(ctx_dft_dec), 0, 1);
628+
629+
common_params_sampling params;
630+
params.no_perf = false;
631+
params.top_k = 1;
632+
params.samplers = { COMMON_SAMPLER_TYPE_TOP_K };
633+
smpl = common_sampler_init(llama_get_model(ctx_dft_dec), params);
634+
}
635+
636+
~common_speculative_state_dflash() override {
637+
llama_perf_context_print(ctx_dft_dec);
638+
639+
if (ctx_dft_dec) {
640+
llama_free(ctx_dft_dec);
641+
}
642+
643+
if (ctx_dft_enc) {
644+
llama_free(ctx_dft_enc);
645+
}
646+
647+
common_sampler_free(smpl);
648+
llama_batch_free(batch);
649+
}
650+
651+
void begin(const llama_tokens & prompt) override {
652+
GGML_UNUSED(prompt);
653+
}
654+
655+
void draft(
656+
const common_params_speculative & params,
657+
const llama_tokens & prompt_tgt,
658+
llama_token id_last,
659+
llama_tokens & result) override {
660+
const int n_embd = llama_model_n_embd(llama_get_model(ctx_dft_dec));
661+
// block_size is bounded by the model's trained block_size (from GGUF metadata).
662+
const int model_block_size = llama_model_dflash_block_size(llama_get_model(ctx_dft_dec));
663+
const int block_size = std::min((int)params.n_max, model_block_size);
664+
const int n = (int)prompt_tgt.size();
665+
const int n_new = n - dflash_n_past;
666+
667+
GGML_ASSERT(n >= 1 && "prompt_tgt is empty");
668+
GGML_ASSERT(n_new >= 1 && "must have at least 1 new token");
669+
670+
// Step 1: Encode new accepted tokens' features
671+
const float * features = llama_get_dflash_target_features(ctx_tgt);
672+
673+
llama_batch enc_batch = {
674+
/*.n_tokens =*/ n_new,
675+
/*.token =*/ nullptr,
676+
/*.embd =*/ const_cast<float*>(features),
677+
/*.pos =*/ nullptr,
678+
/*.n_seq_id =*/ nullptr,
679+
/*.seq_id =*/ nullptr,
680+
/*.logits =*/ nullptr,
681+
};
682+
if (llama_encode(ctx_dft_enc, enc_batch) != 0) {
683+
LOG_ERR("DFlash: encoder failed\n");
684+
return;
685+
}
686+
687+
const float * target_ctx_new = llama_get_embeddings(ctx_dft_enc);
688+
GGML_ASSERT(target_ctx_new && "encoder output is null");
689+
690+
// Step 2: Append to accumulated target_ctx and set on decoder context (writes to cross.v_embd)
691+
const size_t new_size = (size_t)n_embd * n_new;
692+
accumulated_ctx.insert(accumulated_ctx.end(), target_ctx_new, target_ctx_new + new_size);
693+
694+
const int n_ctx_total = (int)(accumulated_ctx.size() / n_embd);
695+
llama_set_dflash_accumulated_target_ctx(ctx_dft_dec, accumulated_ctx.data(), n_embd, n_ctx_total);
696+
697+
// Step 3: Decode noise block
698+
const llama_token mask_token_id = llama_model_dflash_mask_token_id(llama_get_model(ctx_dft_dec));
699+
700+
common_batch_clear(batch);
701+
for (int i = 0; i < block_size; i++) {
702+
const llama_token tok = (i == 0) ? id_last : mask_token_id;
703+
common_batch_add(batch, tok, i, {0}, true);
704+
}
705+
706+
if (llama_decode(ctx_dft_dec, batch) != 0) {
707+
LOG_ERR("DFlash: noise decode failed\n");
708+
return;
709+
}
710+
711+
dflash_n_past = n;
712+
713+
// Step 4: Sample draft tokens from positions 1..block_size-1
714+
result.clear();
715+
common_sampler_reset(smpl);
716+
717+
for (int i = 1; i < block_size; i++) {
718+
common_sampler_sample(smpl, ctx_dft_dec, i);
719+
720+
const auto * cur_p = common_sampler_get_candidates(smpl, true);
721+
const llama_token id = cur_p->data[0].id;
722+
723+
common_sampler_accept(smpl, id, true);
724+
result.push_back(id);
725+
}
726+
}
727+
728+
void accept(uint16_t n_accepted) override {
729+
GGML_UNUSED(n_accepted);
730+
}
731+
};
732+
598733
// state of self-speculation (simple implementation, not ngram-map)
599734
struct common_speculative_state_ngram_simple : public common_speculative_state {
600735
common_ngram_simple_config config;
@@ -978,13 +1113,13 @@ common_speculative * common_speculative_init(
9781113
llama_context * ctx_dft_dec = nullptr;
9791114

9801115
if (params.model_dft) {
981-
if (params.eagle3) {
1116+
if (params.eagle3 || params.dflash) {
9821117
llama_context_params params_enc = params.cparams_dft;
9831118
params_enc.target_model = nullptr;
9841119
params_enc.embeddings = true;
9851120
ctx_dft_enc = llama_init_from_model(params.model_dft, params_enc);
9861121
if (!ctx_dft_enc) {
987-
LOG_ERR("failed to create EAGLE3 encoder context\n");
1122+
LOG_ERR("failed to create %s draft model encoder context\n", params.eagle3 ? "EAGLE3" : "DFlash");
9881123
return nullptr;
9891124
}
9901125

@@ -993,13 +1128,13 @@ common_speculative * common_speculative_init(
9931128
params_dec.embeddings = true;
9941129
ctx_dft_dec = llama_init_from_model(params.model_dft, params_dec);
9951130
if (!ctx_dft_dec) {
996-
LOG_ERR("failed to create EAGLE3 decoder context\n");
1131+
LOG_ERR("failed to create %s draft model decoder context\n", params.eagle3 ? "EAGLE3" : "DFlash");
9971132
return nullptr;
9981133
}
9991134
} else {
10001135
ctx_dft = llama_init_from_model(params.model_dft, params.cparams_dft);
10011136
if (ctx_dft == nullptr) {
1002-
LOG_ERR("%s", "failed to create draft context\n");
1137+
LOG_ERR("failed to create draft model context\n");
10031138
return nullptr;
10041139
}
10051140
}
@@ -1010,6 +1145,7 @@ common_speculative * common_speculative_init(
10101145
{
10111146
bool has_draft = !params.mparams_dft.path.empty();
10121147
bool has_draft_eagle3 = params.eagle3;
1148+
bool has_draft_dflash = params.dflash;
10131149

10141150
bool has_ngram_cache = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_CACHE);
10151151
bool has_ngram_simple = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE);
@@ -1052,6 +1188,8 @@ common_speculative * common_speculative_init(
10521188
if (has_draft) {
10531189
if (has_draft_eagle3) {
10541190
configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_EAGLE3, params));
1191+
} else if (has_draft_dflash) {
1192+
configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DFLASH, params));
10551193
} else {
10561194
configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT, params));
10571195
}
@@ -1081,6 +1219,14 @@ common_speculative * common_speculative_init(
10811219
));
10821220
break;
10831221
}
1222+
case COMMON_SPECULATIVE_TYPE_DFLASH: {
1223+
impls.push_back(std::make_unique<common_speculative_state_dflash>(config.type,
1224+
/* .ctx_tgt = */ ctx_tgt,
1225+
/* .ctx_dft_enc = */ ctx_dft_enc,
1226+
/* .ctx_dft_dec = */ ctx_dft_dec
1227+
));
1228+
break;
1229+
}
10841230
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: {
10851231
common_ngram_map ngram_map = get_common_ngram_map(config);
10861232

‎convert_hf_to_gguf.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4806,6 +4806,47 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter
48064806
yield from super().modify_tensors(data_torch, name, bid)
48074807

48084808

4809+
@ModelBase.register("DFlashDraftModel")
4810+
class DFlashModel(Qwen3Model):
4811+
model_arch = gguf.MODEL_ARCH.DFLASH
4812+
4813+
def set_vocab(self):
4814+
if self.target_model_dir is None:
4815+
raise ValueError(
4816+
"DFlash draft model requires --target-model-dir to be specified. "
4817+
"Please provide the path to the target model directory containing the tokenizer."
4818+
)
4819+
logger.info(f"DFLASH: Using tokenizer from target model: {self.target_model_dir}")
4820+
original_dir = self.dir_model
4821+
self.dir_model = self.target_model_dir
4822+
super().set_vocab()
4823+
self.dir_model = original_dir
4824+
4825+
def set_gguf_parameters(self):
4826+
super().set_gguf_parameters()
4827+
block_size = self.hparams.get("block_size", 16)
4828+
self.gguf_writer.add_uint32(f"{self.gguf_writer.arch}.block_size", block_size)
4829+
dflash_config = self.hparams.get("dflash_config", {})
4830+
target_layer_ids = dflash_config.get("target_layer_ids", [])
4831+
if target_layer_ids:
4832+
extract_layer_ids = [i + 1 for i in target_layer_ids]
4833+
self.gguf_writer.add_array(f"{self.gguf_writer.arch}.target_layer_ids", extract_layer_ids)
4834+
mask_token_id = dflash_config.get("mask_token_id", None)
4835+
if mask_token_id is not None:
4836+
self.gguf_writer.add_uint32(f"{self.gguf_writer.arch}.mask_token_id", mask_token_id)
4837+
4838+
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
4839+
if name == "fc.weight":
4840+
yield (name, data_torch)
4841+
return
4842+
if name == "hidden_norm.weight":
4843+
yield ("hidden_norm.weight", data_torch)
4844+
return
4845+
if not name.startswith("model."):
4846+
name = "model." + name
4847+
yield from super().modify_tensors(data_torch, name, bid)
4848+
4849+
48094850
@ModelBase.register("Qwen3MoeForCausalLM")
48104851
class Qwen3MoeModel(Qwen2MoeModel):
48114852
model_arch = gguf.MODEL_ARCH.QWEN3MOE

0 commit comments

Comments
 (0)