Skip to content

spec : add DFlash2 support (local convolution + candidate selector) (#27342) #27816

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

Merged
ngxson merged 2 commits into from
Aug 27, 2026
Merged

spec : add DFlash2 support (local convolution + candidate selector) (#27342) #27816

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4a6ad48
spec : add DFlash2 support (local convolution + candidate selector) (…
SubSir Aug 27, 2026
4e2f54f
revert top-k.cu changes
ngxson Aug 27, 2026
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
74 changes: 69 additions & 5 deletions common/speculative.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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());
Expand All @@ -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;
Expand All @@ -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));
Expand All @@ -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
}

Expand Down Expand Up @@ -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,
Expand All @@ -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;
Expand Down Expand Up @@ -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);
}
}

Expand Down Expand Up @@ -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;
Expand Down
1 change: 1 addition & 0 deletions conversion/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
"DeepseekV3ForCausalLM": "deepseek",
"DeepseekV32ForCausalLM": "deepseek",
"DFlashDraftModel": "qwen",
"DFlash2DraftModel": "qwen",
"Qwen3DSparkModel": "qwen",
"DSparkDraftModel": "qwen",
"DSparkSpeculator": "qwen",
Expand Down
62 changes: 59 additions & 3 deletions conversion/qwen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -639,7 +639,7 @@ class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
model_arch = gguf.MODEL_ARCH.QWEN35MOE


@ModelBase.register("DFlashDraftModel")
@ModelBase.register("DFlashDraftModel", "DFlash2DraftModel")
@ModelBase.example("z-lab/Qwen3.5-9B-DFlash")
class DFlashModel(Qwen3Model):
model_arch = gguf.MODEL_ARCH.DFLASH
Expand Down Expand Up @@ -678,9 +678,31 @@ def set_vocab(self):
def set_gguf_parameters(self):
super().set_gguf_parameters()

block_size = self.hparams.get("block_size", 16)
self.gguf_writer.add_block_size(block_size)
dflash_config = self.hparams.get("dflash_config", {})
block_size = dflash_config.get("block_size", self.hparams.get("block_size", 16))
self.gguf_writer.add_block_size(block_size)

if "conv_kernel_size" in dflash_config:
self.gguf_writer.add_conv_kernel_size(int(dflash_config["conv_kernel_size"]))
self.gguf_writer.add_conv_group_size(int(dflash_config["conv_group_size"]))
self.gguf_writer.add_selector_rank(int(dflash_config["selector_rank"]))
self.gguf_writer.add_selector_top_k(int(dflash_config["selector_top_k"]))

output_multiplier = dflash_config.get(
"output_multiplier", self.hparams.get("output_multiplier")
)
if output_multiplier is not None:
self.gguf_writer.add_logit_scale(float(output_multiplier))
softcap = dflash_config.get(
"final_logit_softcapping", self.hparams.get("final_logit_softcapping")
)
if softcap is not None and float(softcap) > 0:
self.gguf_writer.add_final_logit_softcapping(float(softcap))
embedding_scale = dflash_config.get(
"input_embedding_scale", self.hparams.get("input_embedding_scale")
)
if embedding_scale is not None:
self.gguf_writer.add_embedding_scale(float(embedding_scale))

target_layer_ids = dflash_config.get("target_layer_ids", [])
if target_layer_ids:
Expand All @@ -695,17 +717,51 @@ def set_gguf_parameters(self):
self.gguf_writer.add_sliding_window(sliding_window)
self.gguf_writer.add_sliding_window_pattern(is_swa)

# M-RoPE target: the draft ropes on the temporal dim only, so write
# degenerate sections [n_rot/2, 0, 0, 0]
if self._target_uses_mrope():
head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
self.gguf_writer.add_rope_dimension_sections([head_dim // 2, 0, 0, 0])

def _target_uses_mrope(self) -> bool:
if self.target_model_dir is None:
return False
with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f:
cfg = json.load(f)
cfg = cfg.get("text_config", cfg)
rope = cfg.get("rope_parameters") or cfg.get("rope_scaling") or {}
return "mrope_section" in rope

@classmethod
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
name, gen = item
if not name.startswith("model."):
name = "model." + name
return super().filter_tensors((name, gen))

_ROPE_PERMUTE_SUFFIXES = (
"self_attn.q_proj.weight",
"self_attn.k_proj.weight",
"self_attn.q_norm.weight",
"self_attn.k_norm.weight",
)

def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True):
return

# interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd
if not self.hparams.get("rope_is_neox_style", True) and name.endswith(self._ROPE_PERMUTE_SUFFIXES):
head_dim = self.hparams["head_dim"]
shape = data_torch.shape
data_torch = data_torch.reshape(-1, head_dim // 2, 2, *shape[1:]).transpose(1, 2).reshape(shape)

if name in (
"model.candidate_selector.predecessor_codebook",
"model.candidate_selector.successor_codebook",
):
name += ".weight"

yield from super().modify_tensors(data_torch, name, bid)


Expand Down
25 changes: 25 additions & 0 deletions gguf-py/gguf/constants.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,10 @@ class LLM:
TARGET_LAYERS = "{arch}.target_layers"
TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size"
BLOCK_SIZE = "{arch}.block_size"
CONV_KERNEL_SIZE = "{arch}.conv_kernel_size"
CONV_GROUP_SIZE = "{arch}.conv_group_size"
SELECTOR_RANK = "{arch}.selector_rank"
SELECTOR_TOP_K = "{arch}.selector_top_k"
SAMPLE_FROM_ANCHOR = "{arch}.sample_from_anchor"
NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual"
NORM_BEFORE_FC = "{arch}.norm_before_fc"
Expand Down Expand Up @@ -1146,6 +1150,13 @@ class MODEL_TENSOR(IntEnum):
DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed
DSPARK_MARKOV_W2 = auto() # markov head: bias projection
DSPARK_CONF_PROJ = auto() # confidence head
DFLASH_ATTN_CONV_BASE = auto()
DFLASH_ATTN_CONV_PROJ = auto()
DFLASH_FFN_CONV_BASE = auto()
DFLASH_FFN_CONV_PROJ = auto()
DFLASH_SELECTOR_PREV = auto()
DFLASH_SELECTOR_NEXT = auto()
DFLASH_SELECTOR_HIDDEN = auto()
# lfm2 audio
A_ENC_NORM_CONV = auto()
A_ENC_LINEAR_POS = auto()
Expand Down Expand Up @@ -1895,6 +1906,13 @@ class MODEL_TENSOR(IntEnum):
MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1",
MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2",
MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj",
MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base",
MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj",
MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base",
MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: "blk.{bid}.ffn_conv_proj",
MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor",
MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor",
MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden",
MODEL_TENSOR.D2T: "d2t",
}

Expand Down Expand Up @@ -4953,6 +4971,13 @@ class MODEL_TENSOR(IntEnum):
MODEL_TENSOR.DSPARK_MARKOV_W1,
MODEL_TENSOR.DSPARK_MARKOV_W2,
MODEL_TENSOR.DSPARK_CONF_PROJ,
MODEL_TENSOR.DFLASH_ATTN_CONV_BASE,
MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ,
MODEL_TENSOR.DFLASH_FFN_CONV_BASE,
MODEL_TENSOR.DFLASH_FFN_CONV_PROJ,
MODEL_TENSOR.DFLASH_SELECTOR_PREV,
MODEL_TENSOR.DFLASH_SELECTOR_NEXT,
MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN,
],
MODEL_ARCH.MISTRAL4: [
MODEL_TENSOR.TOKEN_EMBD,
Expand Down
12 changes: 12 additions & 0 deletions gguf-py/gguf/gguf_writer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -993,6 +993,18 @@ def add_sliding_window(self, value: int) -> None:
def add_block_size(self, value: int) -> None:
self.add_uint32(Keys.LLM.BLOCK_SIZE.format(arch=self.arch), value)

def add_conv_kernel_size(self, value: int) -> None:
self.add_uint32(Keys.LLM.CONV_KERNEL_SIZE.format(arch=self.arch), value)

def add_conv_group_size(self, value: int) -> None:
self.add_uint32(Keys.LLM.CONV_GROUP_SIZE.format(arch=self.arch), value)

def add_selector_rank(self, value: int) -> None:
self.add_uint32(Keys.LLM.SELECTOR_RANK.format(arch=self.arch), value)

def add_selector_top_k(self, value: int) -> None:
self.add_uint32(Keys.LLM.SELECTOR_TOP_K.format(arch=self.arch), value)

def add_sample_from_anchor(self, value: bool) -> None:
self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.format(arch=self.arch), value)

Expand Down
28 changes: 28 additions & 0 deletions gguf-py/gguf/tensor_mapping.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1355,6 +1355,34 @@ class TensorNameMap:
"model.confidence_head.proj", # dspark
),

MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: (
"model.layers.{bid}.attention_conv.base_kernel",
),

MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: (
"model.layers.{bid}.attention_conv.kernel_projection",
),

MODEL_TENSOR.DFLASH_FFN_CONV_BASE: (
"model.layers.{bid}.mlp_conv.base_kernel",
),

MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: (
"model.layers.{bid}.mlp_conv.kernel_projection",
),

MODEL_TENSOR.DFLASH_SELECTOR_PREV: (
"model.candidate_selector.predecessor_codebook",
),

MODEL_TENSOR.DFLASH_SELECTOR_NEXT: (
"model.candidate_selector.successor_codebook",
),

MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: (
"model.candidate_selector.hidden_projection",
),

MODEL_TENSOR.CLS: (
"classifier", # jina
"classifier.dense", # roberta
Expand Down
20 changes: 20 additions & 0 deletions src/llama-arch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -344,6 +344,12 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" },
{ LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" },

{ LLM_KV_DFLASH_BLOCK_SIZE, "%s.block_size" },
{ LLM_KV_DFLASH_CONV_KERNEL_SIZE, "%s.conv_kernel_size" },
{ LLM_KV_DFLASH_CONV_GROUP_SIZE, "%s.conv_group_size" },
{ LLM_KV_DFLASH_SELECTOR_RANK, "%s.selector_rank" },
{ LLM_KV_DFLASH_SELECTOR_TOP_K, "%s.selector_top_k" },

{ LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" },
// sentence-transformers dense modules feature dims
{ LLM_KV_DENSE_2_FEAT_IN, "%s.dense_2_feat_in" },
Expand Down Expand Up @@ -651,6 +657,13 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" },
{ LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" },
{ LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" },
{ LLM_TENSOR_DFLASH_ATTN_CONV_BASE, "blk.%d.attn_conv_base" },
{ LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "blk.%d.attn_conv_proj" },
{ LLM_TENSOR_DFLASH_FFN_CONV_BASE, "blk.%d.ffn_conv_base" },
{ LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "blk.%d.ffn_conv_proj" },
{ LLM_TENSOR_DFLASH_SELECTOR_PREV, "selector_predecessor" },
{ LLM_TENSOR_DFLASH_SELECTOR_NEXT, "selector_successor" },
{ LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "selector_hidden" },
};

// declare information about the model weight tensors:
Expand Down Expand Up @@ -916,6 +929,13 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
{LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
{LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
{LLM_TENSOR_DFLASH_ATTN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_DFLASH_FFN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_DFLASH_FFN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_DFLASH_SELECTOR_PREV, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
{LLM_TENSOR_DFLASH_SELECTOR_NEXT, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
{LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
};

LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {}
Expand Down
Loading
Loading