From 2863b1c7d7352090f6e1ebef382ee9c6fbe65243 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:19:16 -0700 Subject: [PATCH 1/9] dflash: AngelSpec DFly drafter support DFly extends the DFlash draft path with two things, and this wires both onto the existing multi-target-layer capture that DFlash already had: 1. Per-DRAFT-layer context fusion. The shared projection (context_proj -> fc) now only produces a base context; each draft layer adds its own softmax-weighted mix of the raw per-target-layer features, so the encoder emits one context per draft layer instead of one shared context. The decoder's embd batch slices its own layer's context back out for the K/V injection. DFly drops DFlash's encoder hidden_norm in favour of a post-fusion context_norm, so exactly one of the two is present. 2. A TreeFlash predecessor correction chained across the block: position i's logits come from the draft hidden state at row i corrected by the embedding of the token drafted at row i-1, then projected through the target head. Slot 0 is the committed anchor, not a prediction slot, matching the row convention the DFlash reader in common/speculative.cpp already uses. The whole host-side contract change is n_embd_out(): both the nextn output buffer and the embd batch that feeds it back are already sized from it, so widening it for DFly needs no plumbing changes. Nothing is needed in common/speculative.cpp either -- DFly reports no Markov head, so it takes the DFlash reader (rows 1..n-1) and the plain-probability draft cut, which is the threshold form of D-cut that path already implements. DFly is detected from layer_fusion rather than a KV so a DFly export cannot load as plain DFlash, and a checkpoint declaring both a DFly fusion and a DSpark Markov head is rejected with that reason rather than failing later on a missing tensor. Reduced draft vocabularies (d2t) are rejected: the reference chain runs the full target head. Conversion covers Qwen3DFlyModel and the Qwen3DSparkDFlareV2Model alias, and refuses a config whose target_layer_ids exceed its stated target depth -- the reference checkpoint's published metadata does exactly that, and it loads clean and then mis-projects. tests/test-dfly-fusion.cpp pins the fusion math and the resulting flat layout against an independent scalar reference; the axis handling needs two reshape/permute round trips and would otherwise be easy to get wrong in a way that still loads and still drafts plausibly. tests/gen-tiny-dfly.py builds a tiny random DFly GGUF for loader smoke tests. Not yet measured: acceptance rate or speedup against a real target/drafter pair. --- conversion/qwen.py | 77 +++++++++++ gguf-py/gguf/constants.py | 22 ++++ gguf-py/gguf/tensor_mapping.py | 34 ++++- src/llama-arch.cpp | 14 ++ src/llama-arch.h | 7 + src/llama-model.h | 10 ++ src/models/dflash.cpp | 227 +++++++++++++++++++++++++++++++-- tests/CMakeLists.txt | 1 + tests/gen-tiny-dfly.py | 120 +++++++++++++++++ tests/test-dfly-fusion.cpp | 172 +++++++++++++++++++++++++ 10 files changed, 673 insertions(+), 11 deletions(-) create mode 100644 tests/gen-tiny-dfly.py create mode 100644 tests/test-dfly-fusion.cpp diff --git a/conversion/qwen.py b/conversion/qwen.py index ff72c61314f0..dde22ca348cd 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -729,6 +729,83 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield from super().modify_tensors(data_torch, name, bid) +@ModelBase.register("Qwen3DFlyModel", "Qwen3DSparkDFlareV2Model") +@ModelBase.example("AngelSlim/Qwen3-8B-DFly-Block8") +class DFlyModel(DFlashModel): + # AngelSpec DFly = DFlash + (a) a per-DRAFT-layer fusion of the raw target-layer features + # on top of the shared context projection, and (b) a TreeFlash predecessor correction + # applied before the target head. There is no Markov/confidence head, so the runtime reads + # it with the DFlash draft reader (block rows 1..n-1), not the DSpark one. + model_arch = gguf.MODEL_ARCH.DFLASH + + def __init__(self, dir_model, *args, **kwargs): + hparams = kwargs.pop("hparams", None) + if hparams is None: + hparams = ModelBase.load_hparams(dir_model, False) + + # DFly carries target_layer_ids/mask_token_id flat; normalize to DFlash's nested schema + hparams.setdefault("dflash_config", { + k: hparams[k] for k in ("target_layer_ids", "mask_token_id") if k in hparams + }) + + super().__init__(dir_model, *args, hparams=hparams, **kwargs) + + hp = self.hparams + + # The published main-branch config of the reference checkpoint states a target depth and + # vocab that do not match the weights (80/120832 against the real 36/151936), which loads + # clean and then mis-projects. Refuse instead of converting a checkpoint that lies. + target_layers = hp.get("target_num_hidden_layers") + layer_ids = hp.get("target_layer_ids") or [] + if target_layers and layer_ids and max(layer_ids) >= int(target_layers): + raise ValueError( + f"DFly target_layer_ids {layer_ids} exceed target_num_hidden_layers {target_layers}. " + "The config metadata does not describe the weights -- pin a known-good revision " + "(the reference checkpoint's is 5712926)." + ) + + if int(hp.get("target_hidden_size", hp["hidden_size"])) != int(hp["hidden_size"]): + raise ValueError( + "DFly residual fusion requires target_hidden_size == hidden_size, got " + f"{hp.get('target_hidden_size')} vs {hp['hidden_size']}." + ) + + if hp.get("markov_rank") or hp.get("enable_confidence_head"): + raise ValueError( + "DFly does not use the DSpark Markov/confidence head, but this config declares one. " + "A drafter reporting a Markov head is read one block row late by the runtime." + ) + + self._has_correction = bool(hp.get("enable_hidden_correction", True)) + if self._has_correction: + correction_type = hp.get("hidden_correction_type", "swiglu") + if correction_type != "swiglu": + raise ValueError(f"unsupported hidden_correction_type {correction_type!r} (only 'swiglu')") + + # slot 0 of a DFly block is the committed bonus anchor, not a prediction slot + self._sample_from_anchor = not bool(hp.get("dspark_bonus_anchor", True)) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor) + + def prepare_tensors(self): + super().prepare_tensors() + if self._has_correction and not self._seen_correction: + raise ValueError( + "config sets enable_hidden_correction but no hidden_correction.* weights were " + "found; the export is incomplete and would draft without the correction." + ) + + _seen_correction = False + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if "hidden_correction." in name: + self._seen_correction = True + + yield from super().modify_tensors(data_torch, name, bid) + + @ModelBase.register( "Qwen3DSparkModel", "DSparkDraftModel", diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d2284150ab6b..dff3268553b2 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -1146,6 +1146,14 @@ 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 + # dfly + DFLY_LAYER_FUSION = auto() # per-draft-layer context mixing logits + DFLY_CTX_NORM = auto() # post-fusion context norm + DFLY_HC_HIDDEN_NORM = auto() # predecessor correction, hidden branch + DFLY_HC_EMBED_NORM = auto() # predecessor correction, embedding branch + DFLY_HC_GATE = auto() + DFLY_HC_UP = auto() + DFLY_HC_DOWN = auto() # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -1893,6 +1901,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", MODEL_TENSOR.FC: "fc", MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", + MODEL_TENSOR.DFLY_LAYER_FUSION: "layer_fusion", + MODEL_TENSOR.DFLY_CTX_NORM: "context_norm", + MODEL_TENSOR.DFLY_HC_HIDDEN_NORM: "hidden_correction.hidden_norm", + MODEL_TENSOR.DFLY_HC_EMBED_NORM: "hidden_correction.embed_norm", + MODEL_TENSOR.DFLY_HC_GATE: "hidden_correction.gate", + MODEL_TENSOR.DFLY_HC_UP: "hidden_correction.up", + MODEL_TENSOR.DFLY_HC_DOWN: "hidden_correction.down", MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", MODEL_TENSOR.D2T: "d2t", @@ -4951,6 +4966,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.D2T, # optional DSpark heads MODEL_TENSOR.DSPARK_MARKOV_W1, + MODEL_TENSOR.DFLY_LAYER_FUSION, + MODEL_TENSOR.DFLY_CTX_NORM, + MODEL_TENSOR.DFLY_HC_HIDDEN_NORM, + MODEL_TENSOR.DFLY_HC_EMBED_NORM, + MODEL_TENSOR.DFLY_HC_GATE, + MODEL_TENSOR.DFLY_HC_UP, + MODEL_TENSOR.DFLY_HC_DOWN, MODEL_TENSOR.DSPARK_MARKOV_W2, MODEL_TENSOR.DSPARK_CONF_PROJ, ], diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index ef580518e97d..5f6e62d33e09 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -95,6 +95,7 @@ class TensorNameMap: ), # Output norm MODEL_TENSOR.OUTPUT_NORM: ( + "model.final_norm", # dfly "gpt_neox.final_layer_norm", # gptneox "transformer.ln_f", # gpt2 gpt-j falcon jais exaone "model.norm", # llama-hf baichuan internlm2 olmoe olmo2 phimoe plamo2 @@ -1339,8 +1340,37 @@ class TensorNameMap: ), MODEL_TENSOR.FC: ( - "model.fc", # dflash - "encoder.fc", # dflash (transformers MuseGlimmerAssistant) + "model.fc", # dflash + "encoder.fc", # dflash (transformers MuseGlimmerAssistant) + "model.context_proj", # dfly (the shared base context projection) + ), + + MODEL_TENSOR.DFLY_LAYER_FUSION: ( + "model.layer_fusion_weights", # dfly + ), + + MODEL_TENSOR.DFLY_CTX_NORM: ( + "model.context_norm", # dfly + ), + + MODEL_TENSOR.DFLY_HC_HIDDEN_NORM: ( + "model.hidden_correction.hidden_norm", # dfly + ), + + MODEL_TENSOR.DFLY_HC_EMBED_NORM: ( + "model.hidden_correction.embed_norm", # dfly + ), + + MODEL_TENSOR.DFLY_HC_GATE: ( + "model.hidden_correction.gate_proj", # dfly + ), + + MODEL_TENSOR.DFLY_HC_UP: ( + "model.hidden_correction.up_proj", # dfly + ), + + MODEL_TENSOR.DFLY_HC_DOWN: ( + "model.hidden_correction.down_proj", # dfly ), MODEL_TENSOR.DSPARK_MARKOV_W1: ( diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 8e0b369e2e63..895e47a24576 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -657,6 +657,13 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" }, { LLM_TENSOR_DSPARK_LOG_SNR_FC1, "log_snr_fc1" }, { LLM_TENSOR_DSPARK_LOG_SNR_FC2, "log_snr_fc2" }, + { LLM_TENSOR_DFLY_LAYER_FUSION, "layer_fusion" }, + { LLM_TENSOR_DFLY_CTX_NORM, "context_norm" }, + { LLM_TENSOR_DFLY_HC_HIDDEN_NORM, "hidden_correction.hidden_norm" }, + { LLM_TENSOR_DFLY_HC_EMBED_NORM, "hidden_correction.embed_norm" }, + { LLM_TENSOR_DFLY_HC_GATE, "hidden_correction.gate" }, + { LLM_TENSOR_DFLY_HC_UP, "hidden_correction.up" }, + { LLM_TENSOR_DFLY_HC_DOWN, "hidden_correction.down" }, }; // declare information about the model weight tensors: @@ -920,6 +927,13 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_D2T, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, // dspark {LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_DFLY_LAYER_FUSION, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLY_CTX_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_DFLY_HC_HIDDEN_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_DFLY_HC_EMBED_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_DFLY_HC_GATE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLY_HC_UP, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLY_HC_DOWN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {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_DSPARK_LOG_SNR_FC1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index 2b7172e8f820..db4bf5a1b92d 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -665,6 +665,13 @@ enum llm_tensor { LLM_TENSOR_DSPARK_CONF_PROJ, LLM_TENSOR_DSPARK_LOG_SNR_FC1, LLM_TENSOR_DSPARK_LOG_SNR_FC2, + LLM_TENSOR_DFLY_LAYER_FUSION, + LLM_TENSOR_DFLY_CTX_NORM, + LLM_TENSOR_DFLY_HC_HIDDEN_NORM, + LLM_TENSOR_DFLY_HC_EMBED_NORM, + LLM_TENSOR_DFLY_HC_GATE, + LLM_TENSOR_DFLY_HC_UP, + LLM_TENSOR_DFLY_HC_DOWN, }; diff --git a/src/llama-model.h b/src/llama-model.h index ff1bc9b563e3..4d56f8cbdb65 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -655,6 +655,16 @@ struct llama_model { struct ggml_tensor * dspark_log_snr_fc2_w = nullptr; // [n_embd -> n_embd] struct ggml_tensor * dspark_log_snr_fc2_b = nullptr; + // AngelSpec DFly: per-draft-layer target-context fusion + TreeFlash predecessor correction. + // dfly_layer_fusion is the discriminant: present => DFly, absent => plain DFlash/DSpark. + struct ggml_tensor * dfly_layer_fusion = nullptr; // [n_ctx_feat, n_layer] fusion logits + struct ggml_tensor * dfly_ctx_norm = nullptr; // post-fusion context norm (replaces output_norm_enc) + struct ggml_tensor * dfly_hc_hidden_norm = nullptr; + struct ggml_tensor * dfly_hc_embed_norm = nullptr; + struct ggml_tensor * dfly_hc_gate = nullptr; // [2*n_embd, n_ff_hc] + struct ggml_tensor * dfly_hc_up = nullptr; // [2*n_embd, n_ff_hc] + struct ggml_tensor * dfly_hc_down = nullptr; // [n_ff_hc, n_embd] + // unified vector to store target-model extracted layer ids in eagle3, dflash, etc. std::vector target_layer_ids; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 8f98f332b420..9752be48ffca 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -16,6 +16,17 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { hparams.n_embd_inp_enc_impl = (uint32_t) target_layer_ids.size() * hparams.n_embd; + // AngelSpec DFly fuses the target context once per DRAFT layer, so the encoder emits + // [n_embd, n_layer] per token instead of a single [n_embd] row. Both the nextn output + // buffer and the embd batch that feeds it back are sized from n_embd_out(), so widening + // it here is the whole host-side contract change. Detected from the fusion tensor rather + // than a KV so a DFly export cannot load as plain DFlash. + if (ml.get_tensor_meta("layer_fusion.weight")) { + hparams.n_embd_out_impl = hparams.n_layer() * hparams.n_embd; + LLAMA_LOG_INFO("%s: DFly per-layer context fusion (n_layer = %u, n_embd_out = %u)\n", + __func__, hparams.n_layer(), hparams.n_embd_out()); + } + // dspark GIDD log-SNR conditioning (drafters trained with the GIDD bundle); // absent on every other drafter, so it must default off ml.get_key(LLM_KV_LOG_SNR_CONDITIONING, hparams.dspark_log_snr_conditioning, false); @@ -114,6 +125,16 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { LLAMA_LOG_INFO("%s: DFlash using d2t mapping (draft_vocab_size = %lld)\n", __func__, (long long) n_vocab_draft); } + // AngelSpec DFly: per-draft-layer context fusion + TreeFlash predecessor correction. + // Detected from the fusion tensor, and checked before the Markov head below so a + // checkpoint carrying both is reported as such rather than as a missing tensor. + const bool is_dfly = ml->get_tensor_meta("layer_fusion.weight") != nullptr; + + if (is_dfly && d2t) { + throw std::runtime_error("dflash: DFly with a reduced draft vocabulary (d2t) is not supported. " + "The reference chain runs the full target head"); + } + // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head // // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) @@ -129,6 +150,11 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { "The export is incomplete; it would load as plain DFlash and read drafts one row late"); } + if (markov_meta && is_dfly) { + throw std::runtime_error("dflash: checkpoint has both a DFly layer_fusion and a DSpark markov_w1. " + "The two draft chains are mutually exclusive; the export is wrong"); + } + if (markov_meta) { const int64_t dspark_markov_rank = markov_meta->ne[0]; @@ -157,7 +183,35 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED); - output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) + + // DFly replaces the encoder's single hidden_norm with a post-fusion context_norm, so + // exactly one of the two is present. + if (is_dfly) { + const int64_t n_ctx_feat = (int64_t) target_layer_ids.size(); + + dfly_layer_fusion = create_tensor(tn(LLM_TENSOR_DFLY_LAYER_FUSION, "weight"), { n_ctx_feat, n_layer }, 0); + dfly_ctx_norm = create_tensor(tn(LLM_TENSOR_DFLY_CTX_NORM, "weight"), { n_embd }, 0); + + // TreeFlash predecessor correction. Optional in the reference + // (enable_hidden_correction), so absence is a valid checkpoint, but a partial + // set is a broken export rather than something to degrade past. + const struct ggml_tensor * hc_meta = ml->get_tensor_meta("hidden_correction.down.weight"); + if (hc_meta) { + const int64_t n_ff_hc = hc_meta->ne[0]; + + dfly_hc_hidden_norm = create_tensor(tn(LLM_TENSOR_DFLY_HC_HIDDEN_NORM, "weight"), { n_embd }, 0); + dfly_hc_embed_norm = create_tensor(tn(LLM_TENSOR_DFLY_HC_EMBED_NORM, "weight"), { n_embd }, 0); + dfly_hc_gate = create_tensor(tn(LLM_TENSOR_DFLY_HC_GATE, "weight"), { 2*n_embd, n_ff_hc }, 0); + dfly_hc_up = create_tensor(tn(LLM_TENSOR_DFLY_HC_UP, "weight"), { 2*n_embd, n_ff_hc }, 0); + dfly_hc_down = create_tensor(tn(LLM_TENSOR_DFLY_HC_DOWN, "weight"), { n_ff_hc, n_embd }, 0); + + LLAMA_LOG_INFO("%s: DFly predecessor correction (n_ff = %lld)\n", __func__, (long long) n_ff_hc); + } else { + LLAMA_LOG_INFO("%s: DFly without predecessor correction\n", __func__); + } + } else { + output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) + } output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm // optional: reduced-vocab drafts ship their own lm head, full-vocab drafts can share the target's via ctx_other @@ -267,13 +321,46 @@ ggml_tensor * llama_model_dflash::graph::build_inp_embd_enc() const { // DFlash Encoder: processes target model features through feature fusion layer template <> llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { - ggml_tensor * cur = build_inp_embd_enc(); + ggml_tensor * inp = build_inp_embd_enc(); - cur = build_lora_mm(model.fc, cur, model.fc_s); + ggml_tensor * cur = build_lora_mm(model.fc, inp, model.fc_s); cb(cur, "fc_out", -1); - cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1); - cb(cur, "enc_norm_out", -1); + if (model.dfly_layer_fusion) { + // DFly: the shared projection is only the base context. Each draft layer adds its own + // softmax-weighted mix of the raw per-target-layer features, so the encoder emits one + // context per draft layer. Reference: Qwen3DFlyModel.project_target_hidden. + const int64_t n_feat = model.dfly_layer_fusion->ne[0]; + const int64_t n_lyr = model.dfly_layer_fusion->ne[1]; + + GGML_ASSERT(n_feat*n_embd == (int64_t) hparams.n_embd_inp_enc()); + GGML_ASSERT(n_lyr == n_layer); + + // softmax over each draft layer's n_feat mixing logits (torch: softmax(dim=-1) on [L, T]) + ggml_tensor * probs = ggml_soft_max(ctx0, model.dfly_layer_fusion); // [n_feat, n_layer] + + // [n_feat*n_embd, n_tokens] -> [n_feat, n_embd, n_tokens]: put the contracted axis first + ggml_tensor * feats = ggml_cont(ctx0, ggml_permute(ctx0, + ggml_reshape_3d(ctx0, inp, n_embd, n_feat, n_tokens), 1, 0, 2, 3)); + + ggml_tensor * resid = ggml_mul_mat(ctx0, probs, + ggml_reshape_2d(ctx0, feats, n_feat, n_embd*n_tokens)); // [n_layer, n_embd*n_tokens] + + resid = ggml_cont(ctx0, ggml_permute(ctx0, + ggml_reshape_3d(ctx0, resid, n_lyr, n_embd, n_tokens), 1, 0, 2, 3)); // [n_embd, n_layer, n_tokens] + + // broadcast the base context across draft layers + cur = ggml_add(ctx0, resid, ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens)); + cur = build_norm(cur, model.dfly_ctx_norm, NULL, LLM_NORM_RMS, -1); + + // flatten layer-major within each token: the host round-trips this as one + // n_embd_out()-wide row per token and the decoder views layer il back out of it + cur = ggml_reshape_2d(ctx0, cur, n_embd*n_lyr, n_tokens); + cb(cur, "dfly_ctx_out", -1); + } else { + cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1); + cb(cur, "enc_norm_out", -1); + } ggml_set_output(cur); res->t_h_nextn = cur; @@ -389,6 +476,110 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_build_forward_expand(g.gf, out); } +// DFly (AngelSpec): TreeFlash predecessor correction chained across a block. +// +// Position i's logits come from the draft hidden state at row i corrected by the embedding +// of the token drafted at row i-1, then projected through the target head. Slot 0 of each +// block is the committed anchor, not a prediction slot, so the chain seeds from the anchor +// TOKEN and runs i = 1..block_drafts-1 -- the same row convention the DFlash reader in +// common/speculative.cpp uses (rows 1..n-1), which is why a DFly drafter must NOT report a +// DSpark Markov head. +// +// Reference: _DflyDraftSampler.__call__ + DFlyHiddenStatesCorrection.forward. +static void build_dfly_correction_head(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) { + ggml_context * ctx0 = g.ctx0; + auto & res = g.res; + + GGML_ASSERT(model.dfly_hc_hidden_norm && model.dfly_hc_embed_norm && + model.dfly_hc_gate && model.dfly_hc_up && model.dfly_hc_down && + "DFly predecessor-correction weights not loaded"); + + ggml_tensor * hidden = res->t_embd; // [n_embd, n_tokens], after the decoder's final norm + ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens], uncorrected + + const int64_t n_embd = hidden->ne[0]; + const int64_t n_tok = base->ne[1]; + const int64_t n_vocab = base->ne[0]; + + const auto it = model.gguf_kv.find("dflash.block_size"); + GGML_ASSERT(it != model.gguf_kv.end() && "DFly draft requires 'dflash.block_size' in GGUF metadata"); + const int64_t block_size = std::stoi(it->second); + GGML_ASSERT(block_size > 0); + + const int64_t n_blocks = g.ubatch.n_seqs_unq; + GGML_ASSERT(n_blocks > 0 && n_tok % n_blocks == 0 && "DFly head requires equal-size blocks"); + const int64_t block_drafts = n_tok / n_blocks; + if (block_drafts > block_size) { + return; + } + + // the target's head and embeddings when the draft ships none (shared via ctx_other) + auto * output = model.output; + auto * output_s = model.output_s; + auto * tok_embd = model.tok_embd; + if (output == nullptr || tok_embd == nullptr) { + GGML_ASSERT(g.cparams.ctx_other != nullptr); + const auto * model_other = llama_get_model(g.cparams.ctx_other); + if (output == nullptr) { + GGML_ASSERT(model_other->output != nullptr && "DFly head requires the target output projection"); + output = model_other->output; + output_s = model_other->output_s; + } + if (tok_embd == nullptr) { + GGML_ASSERT(model_other->tok_embd != nullptr && "DFly head requires the target token embeddings"); + tok_embd = model_other->tok_embd; + } + } + + const size_t token_stride = (size_t) block_drafts * tokens->nb[0]; + const size_t hidden_stride = (size_t) block_drafts * hidden->nb[1]; + const size_t base_stride = (size_t) block_drafts * base->nb[1]; + + // anchor (committed) token of every block: token 0 of each block, a strided view + ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0); + prev = ggml_cont_1d(ctx0, prev, n_blocks); + + // the anchor slot is not predicted: pass its uncorrected logits through so the output + // keeps one row per ubatch token + ggml_tensor * cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0)); + + for (int64_t i = 1; i < block_drafts; ++i) { + // predecessor correction: delta = down(silu(gate(z)) * up(z)), + // z = [rms(h_i); rms(embd(prev))] + ggml_tensor * prev_embd = ggml_get_rows(ctx0, tok_embd, prev); // [n_embd, n_blocks] + + ggml_tensor * h_i = ggml_cont(ctx0, ggml_view_2d(ctx0, hidden, n_embd, n_blocks, + hidden_stride, i*hidden->nb[1])); + + ggml_tensor * z = ggml_concat(ctx0, + g.build_norm(h_i, model.dfly_hc_hidden_norm, NULL, LLM_NORM_RMS, -1), + g.build_norm(prev_embd, model.dfly_hc_embed_norm, NULL, LLM_NORM_RMS, -1), 0); + + ggml_tensor * delta = g.build_ffn(z, + model.dfly_hc_up, NULL, NULL, + model.dfly_hc_gate, NULL, NULL, + model.dfly_hc_down, NULL, NULL, + NULL, LLM_FFN_SILU, LLM_FFN_PAR, -1); + + ggml_tensor * col = g.build_lora_mm(output, ggml_add(ctx0, h_i, delta), output_s); + + cat = ggml_concat(ctx0, cat, col, 1); + + // greedy chain: the next position conditions on this position's drafted token + if (i + 1 < block_drafts) { + prev = ggml_argmax(ctx0, col); + } + } + + // cat is position-major; restore ubatch block-major order + ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, block_drafts); + out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks] + out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); + + res->t_logits = out; + ggml_build_forward_expand(g.gf, out); +} + // DFlash decoder, dual-mode by batch type: // * embd batch -> fused target features: project + inject K/V into the cache. // * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens @@ -415,9 +606,13 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra // KV cache injection if (ubatch.embd) { - auto inp = std::make_unique(n_embd); + // DFly ships one fused context per draft layer, so the incoming row is n_layer wide + const bool is_dfly = model.dfly_layer_fusion != nullptr; + const int64_t n_embd_batch = is_dfly ? (int64_t) hparams.n_embd_out() : n_embd; - inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens); + auto inp = std::make_unique(n_embd_batch); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_batch, n_tokens); ggml_set_input(inp->embd); ggml_tensor * inp_g = inp->embd; @@ -428,8 +623,17 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra for (int il = 0; il < n_layer; ++il) { const auto & layer = model.layers[il]; - ggml_tensor * Kcur = build_lora_mm(layer.wk, inp_g); - ggml_tensor * Vcur = build_lora_mm(layer.wv, inp_g); + // plain DFlash injects one shared context into every layer; DFly slices out this + // layer's own fused context (encoder writes them layer-major within each token) + ggml_tensor * ctx_il = inp_g; + if (is_dfly) { + ctx_il = ggml_cont(ctx0, ggml_view_2d(ctx0, inp_g, n_embd, n_tokens, + inp_g->nb[1], (size_t) il*n_embd*inp_g->nb[0])); + cb(ctx_il, "dfly_ctx_layer", il); + } + + ggml_tensor * Kcur = build_lora_mm(layer.wk, ctx_il); + ggml_tensor * Vcur = build_lora_mm(layer.wv, ctx_il); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); @@ -653,6 +857,11 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra if (model.dspark_markov_w1) { build_dspark_markov_head(*this, model, inp_tokens); } + + // DFly: re-derive the draft logits through the chained predecessor correction + if (model.dfly_hc_down) { + build_dfly_correction_head(*this, model, inp_tokens); + } } // DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above): diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c56378606bbb..931205179e08 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -159,6 +159,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-grammar-integration.cpp) llama_build_and_test(test-llama-grammar.cpp) llama_build_and_test(test-batch-alloc.cpp) + llama_build_and_test(test-dfly-fusion.cpp) llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) target_link_libraries(test-chat PRIVATE server-context) diff --git a/tests/gen-tiny-dfly.py b/tests/gen-tiny-dfly.py new file mode 100644 index 000000000000..e8f2a567a223 --- /dev/null +++ b/tests/gen-tiny-dfly.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +"""Generate a tiny random DFly (AngelSpec) draft GGUF for loader smoke tests. + +Shapes mirror the reference AngelSlim/Qwen3-8B-DFly-Block8 topology (5 target capture +layers, per-draft-layer context fusion, swiglu predecessor correction) at toy width, with +n_layer deliberately != n_target so a transposed fusion axis cannot pass unnoticed. + + python3 tests/gen-tiny-dfly.py models/ggml-vocab-qwen2.gguf /tmp/tiny-dfly.gguf +""" +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).parent.parent / "gguf-py")) +import gguf # noqa: E402 + +VOCAB_KEY_PREFIXES = ("tokenizer.",) + + +def copy_vocab(writer: gguf.GGUFWriter, vocab_gguf: Path) -> int: + reader = gguf.GGUFReader(vocab_gguf) + n_vocab = 0 + for field in reader.fields.values(): + if not field.name.startswith(VOCAB_KEY_PREFIXES): + continue + if field.types[:1] == [gguf.GGUFValueType.ARRAY]: + sub = field.types[1] + if sub == gguf.GGUFValueType.STRING: + vals = [bytes(field.parts[i]).decode("utf-8") for i in field.data] + writer.add_array(field.name, vals) + if field.name == "tokenizer.ggml.tokens": + n_vocab = len(vals) + else: + vals = [field.parts[i].tolist()[0] for i in field.data] + writer.add_array(field.name, vals) + else: + val = field.parts[field.data[0]] + if field.types[0] == gguf.GGUFValueType.STRING: + writer.add_string(field.name, bytes(val).decode("utf-8")) + else: + writer.add_uint32(field.name, int(val.tolist()[0])) + if n_vocab == 0: + raise SystemExit(f"no tokenizer.ggml.tokens found in {vocab_gguf}") + return n_vocab + + +def main() -> None: + if len(sys.argv) != 3: + raise SystemExit(f"usage: {sys.argv[0]} ") + vocab_gguf, out_path = Path(sys.argv[1]), Path(sys.argv[2]) + + n_embd, n_layer, n_head, n_head_kv = 64, 3, 4, 2 + n_ff, n_feat, n_ff_hc = 128, 5, 64 + head_dim = n_embd // n_head + eps, block_size = 1e-6, 8 + + writer = gguf.GGUFWriter(out_path, "dflash") + n_vocab = copy_vocab(writer, vocab_gguf) + + writer.add_context_length(512) + writer.add_embedding_length(n_embd) + writer.add_block_count(n_layer) + writer.add_feed_forward_length(n_ff) + writer.add_head_count(n_head) + writer.add_head_count_kv(n_head_kv) + writer.add_key_length(head_dim) + writer.add_value_length(head_dim) + writer.add_layer_norm_rms_eps(eps) + writer.add_rope_freq_base(1000000.0) + writer.add_rope_dimension_count(head_dim) + writer.add_block_size(block_size) + writer.add_sample_from_anchor(False) + # +1: the runtime taps a layer's INPUT (same convention as the DFlash converter) + writer.add_target_layers([i + 1 for i in (1, 9, 17, 25, 33)][:n_feat]) + + rng = np.random.default_rng(7) + + def t(name: str, shape: tuple[int, ...]) -> None: + writer.add_tensor(name, rng.standard_normal(shape).astype(np.float32) * 0.05) + + t("token_embd.weight", (n_vocab, n_embd)) + t("output_norm.weight", (n_embd,)) + + # DFly encoder: shared base projection + per-draft-layer fusion, then context_norm. + # Note there is NO enc.output_norm (hidden_norm): DFly replaces it with context_norm. + t("fc.weight", (n_embd, n_feat * n_embd)) + t("layer_fusion.weight", (n_layer, n_feat)) + t("context_norm.weight", (n_embd,)) + + # TreeFlash predecessor correction (swiglu over [hidden ; prev-token embedding]) + t("hidden_correction.hidden_norm.weight", (n_embd,)) + t("hidden_correction.embed_norm.weight", (n_embd,)) + t("hidden_correction.gate.weight", (n_ff_hc, 2 * n_embd)) + t("hidden_correction.up.weight", (n_ff_hc, 2 * n_embd)) + t("hidden_correction.down.weight", (n_embd, n_ff_hc)) + + for i in range(n_layer): + t(f"blk.{i}.attn_norm.weight", (n_embd,)) + t(f"blk.{i}.attn_q.weight", (n_head * head_dim, n_embd)) + t(f"blk.{i}.attn_k.weight", (n_head_kv * head_dim, n_embd)) + t(f"blk.{i}.attn_v.weight", (n_head_kv * head_dim, n_embd)) + t(f"blk.{i}.attn_output.weight", (n_embd, n_head * head_dim)) + t(f"blk.{i}.attn_q_norm.weight", (head_dim,)) + t(f"blk.{i}.attn_k_norm.weight", (head_dim,)) + t(f"blk.{i}.ffn_norm.weight", (n_embd,)) + t(f"blk.{i}.ffn_gate.weight", (n_ff, n_embd)) + t(f"blk.{i}.ffn_up.weight", (n_ff, n_embd)) + t(f"blk.{i}.ffn_down.weight", (n_embd, n_ff)) + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + print(f"wrote {out_path} (n_vocab={n_vocab}, n_layer={n_layer}, n_feat={n_feat}, " + f"expected n_embd_out={n_layer * n_embd})") + + +if __name__ == "__main__": + main() diff --git a/tests/test-dfly-fusion.cpp b/tests/test-dfly-fusion.cpp new file mode 100644 index 000000000000..d0a594b2b837 --- /dev/null +++ b/tests/test-dfly-fusion.cpp @@ -0,0 +1,172 @@ +// DFly (AngelSpec) encoder context fusion: numerical parity + memory-layout contract. +// +// The draft-time graph in src/models/dflash.cpp mixes the raw per-target-layer features into +// one context PER DRAFT LAYER, which needs two reshape/permute round trips to put the +// contracted axis first. Both the axis handling and the resulting flat layout are easy to get +// subtly wrong in a way that still loads and still produces plausible drafts, so this test +// pins them against an independent scalar reference. +// +// Checks: +// 1. fused context == softmax-weighted mix of the raw features + shared base projection +// 2. the flattened [n_embd*n_layer, n_tokens] output is LAYER-MAJOR within each token, +// i.e. the decoder's ggml_view_2d(offset = il*n_embd, stride = nb[1]) recovers layer il + +#include "ggml.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include + +static void graph_compute(ggml_cgraph * gf, int n_threads) { + std::vector buf; + ggml_cplan plan = ggml_graph_plan(gf, n_threads, nullptr); + if (plan.work_size > 0) { + buf.resize(plan.work_size); + plan.work_data = buf.data(); + } + ggml_graph_compute(gf, &plan); +} + +int main() { + // deliberately non-square and n_feat != n_layer so a transposed axis cannot pass + const int64_t n_embd = 8; + const int64_t n_feat = 5; // target capture layers (DFly Qwen3-8B ships 5) + const int64_t n_layer = 3; // draft layers + const int64_t n_tokens = 4; + const float eps = 1e-6f; + + std::vector mem(64u*1024u*1024u); + ggml_init_params ip = { mem.size(), mem.data(), false }; + ggml_context * ctx = ggml_init(ip); + + ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_feat*n_embd, n_tokens); + ggml_tensor * fc = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_feat*n_embd, n_embd); + ggml_tensor * fusion = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_feat, n_layer); + ggml_tensor * cnorm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); + + srand(1234); + auto fill = [](ggml_tensor * t) { + float * d = (float *) t->data; + for (int64_t i = 0; i < ggml_nelements(t); ++i) { + d[i] = 2.0f*((float) rand()/(float) RAND_MAX) - 1.0f; + } + }; + fill(inp); fill(fc); fill(fusion); fill(cnorm); + + // ---- the graph under test: mirrors llama_model_dflash::graph (DFly branch) ---- + ggml_tensor * base = ggml_mul_mat(ctx, fc, inp); // [n_embd, n_tokens] + + ggml_tensor * probs = ggml_soft_max(ctx, fusion); // [n_feat, n_layer] + + ggml_tensor * feats = ggml_cont(ctx, ggml_permute(ctx, + ggml_reshape_3d(ctx, inp, n_embd, n_feat, n_tokens), 1, 0, 2, 3)); + + ggml_tensor * resid = ggml_mul_mat(ctx, probs, + ggml_reshape_2d(ctx, feats, n_feat, n_embd*n_tokens)); // [n_layer, n_embd*n_tokens] + + resid = ggml_cont(ctx, ggml_permute(ctx, + ggml_reshape_3d(ctx, resid, n_layer, n_embd, n_tokens), 1, 0, 2, 3)); + + ggml_tensor * cur = ggml_add(ctx, resid, ggml_reshape_3d(ctx, base, n_embd, 1, n_tokens)); + cur = ggml_mul(ctx, ggml_rms_norm(ctx, cur, eps), cnorm); + cur = ggml_reshape_2d(ctx, cur, n_embd*n_layer, n_tokens); + + // the decoder's per-layer slice of the round-tripped row (layer 1, arbitrary interior pick) + const int64_t il_probe = 1; + ggml_tensor * slice = ggml_cont(ctx, ggml_view_2d(ctx, cur, n_embd, n_tokens, + cur->nb[1], (size_t) il_probe*n_embd*cur->nb[0])); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, cur); + ggml_build_forward_expand(gf, slice); + graph_compute(gf, 2); + + // ---- independent scalar reference ---- + const float * pinp = (const float *) inp->data; + const float * pfc = (const float *) fc->data; + const float * pfus = (const float *) fusion->data; + const float * pcn = (const float *) cnorm->data; + + std::vector ref((size_t) n_embd*n_layer*n_tokens); + + for (int64_t n = 0; n < n_tokens; ++n) { + // shared base projection: base[h] = sum_r fc[r,h] * inp[r,n] + std::vector b(n_embd, 0.0f); + for (int64_t h = 0; h < n_embd; ++h) { + for (int64_t r = 0; r < n_feat*n_embd; ++r) { + b[h] += pfc[r + h*n_feat*n_embd] * pinp[r + n*n_feat*n_embd]; + } + } + + for (int64_t l = 0; l < n_layer; ++l) { + // softmax over this draft layer's n_feat mixing logits + float mx = -INFINITY; + for (int64_t t = 0; t < n_feat; ++t) mx = std::fmax(mx, pfus[t + l*n_feat]); + float sum = 0.0f; + std::vector w(n_feat); + for (int64_t t = 0; t < n_feat; ++t) { w[t] = std::exp(pfus[t + l*n_feat] - mx); sum += w[t]; } + for (int64_t t = 0; t < n_feat; ++t) w[t] /= sum; + + // residual mix of the raw features, then base + residual + std::vector v(n_embd); + for (int64_t h = 0; h < n_embd; ++h) { + float acc = 0.0f; + for (int64_t t = 0; t < n_feat; ++t) { + acc += w[t] * pinp[(h + t*n_embd) + n*n_feat*n_embd]; + } + v[h] = b[h] + acc; + } + + // RMS norm over n_embd, scaled by context_norm + float ss = 0.0f; + for (int64_t h = 0; h < n_embd; ++h) ss += v[h]*v[h]; + const float scale = 1.0f/std::sqrt(ss/(float) n_embd + eps); + for (int64_t h = 0; h < n_embd; ++h) { + ref[(size_t) (h + l*n_embd) + (size_t) n*n_embd*n_layer] = v[h]*scale*pcn[h]; + } + } + } + + // ---- compare ---- + const float * got = (const float *) cur->data; + double max_err = 0.0; + for (size_t i = 0; i < ref.size(); ++i) { + max_err = std::fmax(max_err, std::fabs((double) got[i] - (double) ref[i])); + } + printf("fused context: max abs err = %.3e\n", max_err); + + // layer-major layout: the decoder's strided view must equal reference layer il_probe + const float * pslice = (const float *) slice->data; + double max_err_slice = 0.0; + for (int64_t n = 0; n < n_tokens; ++n) { + for (int64_t h = 0; h < n_embd; ++h) { + const double r = ref[(size_t) (h + il_probe*n_embd) + (size_t) n*n_embd*n_layer]; + max_err_slice = std::fmax(max_err_slice, std::fabs((double) pslice[h + n*n_embd] - r)); + } + } + printf("layer-%lld slice: max abs err = %.3e\n", (long long) il_probe, max_err_slice); + + // a transposed fusion axis would still land within this bound only by coincidence; + // guard against it explicitly by requiring the per-layer contexts to actually differ + double min_sep = INFINITY; + for (int64_t n = 0; n < n_tokens; ++n) { + for (int64_t l = 1; l < n_layer; ++l) { + double d = 0.0; + for (int64_t h = 0; h < n_embd; ++h) { + const double a = got[(size_t) (h) + (size_t) n*n_embd*n_layer]; + const double c = got[(size_t) (h + l*n_embd) + (size_t) n*n_embd*n_layer]; + d += (a-c)*(a-c); + } + min_sep = std::fmin(min_sep, std::sqrt(d)); + } + } + printf("min per-layer context separation = %.3e\n", min_sep); + + ggml_free(ctx); + + const bool ok = max_err < 1e-4 && max_err_slice < 1e-4 && min_sep > 1e-3; + printf("%s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} From 857e5f205e3d5a5b9068325f90c1e0cd2aa24fef Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:56:53 -0700 Subject: [PATCH 2/9] dflash: fix DFly conversion + loading against the reference checkpoint Running the real AngelSlim/Qwen3-8B-DFly-Block8 (revision 5712926) end to end turned up three things the synthetic tests could not: * conversion/__init__.py holds an architecture -> module map used to lazily import converters, separate from @ModelBase.register. Without an entry there the DFly classes never load and conversion fails with "Model Qwen3DFlyModel is not supported". * The fusion parameter is `layer_fusion_weights` in the checkpoint, with no `.weight` suffix, so the writer emits a bare `layer_fusion` tensor. The loader looked for `layer_fusion.weight` and would not have detected DFly at all. Load it suffix-less, like d2t. The tiny-model generator matched the loader rather than the converter, so it agreed with the bug. * The draft context also decodes ordinary batches (context staging before the K/V injection), not just noise blocks. Those have no anchor to condition on: slot 0 is the caller's id_last, which is LLAMA_TOKEN_NULL until the first token is committed. Chaining one of those indexes the embedding table with row -1. Check the block shape before building the chain, and take the anchor embedding from the rows the decoder already gathered instead of re-fetching it by id -- a build-time check alone cannot hold once graphs are reused. Measured with the real pair on an M5 Pro (Qwen3-8B bf16 target + DFly bf16 drafter, Metal): the drafter converts, loads, auto-detects as 'draft-dflash', and runs the fusion and layer-varying injection end to end, producing coherent target output at 12.9 tok/s with draft_n=637. Acceptance is 0, because the predecessor-correction chain is NOT yet enabled. Running it aborts nondeterministically on an uninitialised index (-1) that is not produced by any argmax, in a get_rows outside the chain's own nodes. The chain calls the target's lm_head through ctx_other once per block position and feeds each argmax back into a get_rows on the draft's embeddings; the existing DSpark head never builds such a cross-context dependency, it only biases already-computed logits. That is the suspect and the reason the chain is now opt-in behind LLAMA_DFLY_CHAIN=1, with a load-time warning so a near-zero acceptance run cannot be mistaken for a working DFly. The likely fix is to move the correction out of the model graph into the sampler chain, which is where the reference runtime puts it. --- conversion/__init__.py | 2 ++ src/models/dflash.cpp | 63 ++++++++++++++++++++++++++++++++++-------- tests/gen-tiny-dfly.py | 2 +- 3 files changed, 55 insertions(+), 12 deletions(-) diff --git a/conversion/__init__.py b/conversion/__init__.py index 8de97e95969a..142d42f63f9c 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -54,6 +54,8 @@ "DeepseekV3ForCausalLM": "deepseek", "DeepseekV32ForCausalLM": "deepseek", "DFlashDraftModel": "qwen", + "Qwen3DFlyModel": "qwen", + "Qwen3DSparkDFlareV2Model": "qwen", "Qwen3DSparkModel": "qwen", "DSparkDraftModel": "qwen", "DSparkSpeculator": "qwen", diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 9752be48ffca..687ab4b3cae3 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -21,7 +21,7 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { // buffer and the embd batch that feeds it back are sized from n_embd_out(), so widening // it here is the whole host-side contract change. Detected from the fusion tensor rather // than a KV so a DFly export cannot load as plain DFlash. - if (ml.get_tensor_meta("layer_fusion.weight")) { + if (ml.get_tensor_meta("layer_fusion")) { hparams.n_embd_out_impl = hparams.n_layer() * hparams.n_embd; LLAMA_LOG_INFO("%s: DFly per-layer context fusion (n_layer = %u, n_embd_out = %u)\n", __func__, hparams.n_layer(), hparams.n_embd_out()); @@ -128,7 +128,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { // AngelSpec DFly: per-draft-layer context fusion + TreeFlash predecessor correction. // Detected from the fusion tensor, and checked before the Markov head below so a // checkpoint carrying both is reported as such rather than as a missing tensor. - const bool is_dfly = ml->get_tensor_meta("layer_fusion.weight") != nullptr; + const bool is_dfly = ml->get_tensor_meta("layer_fusion") != nullptr; if (is_dfly && d2t) { throw std::runtime_error("dflash: DFly with a reduced draft vocabulary (d2t) is not supported. " @@ -189,7 +189,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { if (is_dfly) { const int64_t n_ctx_feat = (int64_t) target_layer_ids.size(); - dfly_layer_fusion = create_tensor(tn(LLM_TENSOR_DFLY_LAYER_FUSION, "weight"), { n_ctx_feat, n_layer }, 0); + dfly_layer_fusion = create_tensor(tn(LLM_TENSOR_DFLY_LAYER_FUSION), { n_ctx_feat, n_layer }, 0); dfly_ctx_norm = create_tensor(tn(LLM_TENSOR_DFLY_CTX_NORM, "weight"), { n_embd }, 0); // TreeFlash predecessor correction. Optional in the reference @@ -206,6 +206,10 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { dfly_hc_down = create_tensor(tn(LLM_TENSOR_DFLY_HC_DOWN, "weight"), { n_ff_hc, n_embd }, 0); LLAMA_LOG_INFO("%s: DFly predecessor correction (n_ff = %lld)\n", __func__, (long long) n_ff_hc); + LLAMA_LOG_WARN("%s: DFly predecessor-correction CHAIN IS NOT ENABLED. It is still under " + "development (nondeterministic uninitialised read; see the notes on this branch), " + "so drafts are produced WITHOUT the correction and acceptance will be near zero. " + "Set LLAMA_DFLY_CHAIN=1 to run it anyway.\n", __func__); } else { LLAMA_LOG_INFO("%s: DFly without predecessor correction\n", __func__); } @@ -486,7 +490,8 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & // DSpark Markov head. // // Reference: _DflyDraftSampler.__call__ + DFlyHiddenStatesCorrection.forward. -static void build_dfly_correction_head(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) { +static void build_dfly_correction_head(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens, + ggml_tensor * inp_embd_raw) { ggml_context * ctx0 = g.ctx0; auto & res = g.res; @@ -513,6 +518,29 @@ static void build_dfly_correction_head(llm_graph_context & g, const llama_model return; } + // Only a noise block carries a chain. The draft context also decodes ordinary batches + // (context staging before the K/V injection), and those have no anchor to condition on -- + // their slot 0 is the caller's id_last, which is LLAMA_TOKEN_NULL until the first token is + // committed. Chaining one of those would read token embedding row -1. A block is + // [id_last, MASK, MASK, ...] per sequence, so check that shape before building anything. + if (g.ubatch.token == nullptr) { + return; + } + + const llama_token mask_id = model.vocab.token_mask(); + + for (int64_t b = 0; b < n_blocks; ++b) { + const llama_token anchor = g.ubatch.token[b*block_drafts]; + if (anchor < 0 || anchor >= (llama_token) model.vocab.n_tokens()) { + return; + } + for (int64_t i = 1; i < block_drafts; ++i) { + if (g.ubatch.token[b*block_drafts + i] != mask_id) { + return; + } + } + } + // the target's head and embeddings when the draft ships none (shared via ctx_other) auto * output = model.output; auto * output_s = model.output_s; @@ -531,13 +559,15 @@ static void build_dfly_correction_head(llm_graph_context & g, const llama_model } } - const size_t token_stride = (size_t) block_drafts * tokens->nb[0]; + const size_t hidden_stride = (size_t) block_drafts * hidden->nb[1]; const size_t base_stride = (size_t) block_drafts * base->nb[1]; - // anchor (committed) token of every block: token 0 of each block, a strided view - ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0); - prev = ggml_cont_1d(ctx0, prev, n_blocks); + // position 1 conditions on the block anchor: take its embedding from the rows already + // gathered by the decoder. prev stays null until the chain produces its first token. + ggml_tensor * prev = nullptr; + ggml_tensor * prev_embd = ggml_cont(ctx0, ggml_view_2d(ctx0, inp_embd_raw, n_embd, n_blocks, + (size_t) block_drafts * inp_embd_raw->nb[1], 0)); // the anchor slot is not predicted: pass its uncorrected logits through so the output // keeps one row per ubatch token @@ -546,7 +576,10 @@ static void build_dfly_correction_head(llm_graph_context & g, const llama_model for (int64_t i = 1; i < block_drafts; ++i) { // predecessor correction: delta = down(silu(gate(z)) * up(z)), // z = [rms(h_i); rms(embd(prev))] - ggml_tensor * prev_embd = ggml_get_rows(ctx0, tok_embd, prev); // [n_embd, n_blocks] + if (prev) { + // chained positions condition on the previously drafted token: an argmax, always in range + prev_embd = ggml_get_rows(ctx0, tok_embd, prev); // [n_embd, n_blocks] + } ggml_tensor * h_i = ggml_cont(ctx0, ggml_view_2d(ctx0, hidden, n_embd, n_blocks, hidden_stride, i*hidden->nb[1])); @@ -604,6 +637,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); + // KV cache injection if (ubatch.embd) { // DFly ships one fused context per draft layer, so the incoming row is n_layer wide @@ -703,6 +737,12 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens); cb(inpL, "inp_noise_embd", -1); + // the DFly chain conditions position 1 on the block anchor's embedding; reuse the rows + // gathered here instead of re-fetching by id, so the chain never indexes the embedding + // table with the caller's id_last (which is -1 before the first commit, and which a + // build-time check cannot catch once graphs start being reused across batches) + ggml_tensor * inp_embd_raw = inpL; + // dspark GIDD log-SNR conditioning (LogSnrEmbed): added to the draft noise // embedding before the layer loop, matching the training reference. The // per-position log-SNR is the fixed round-1 inference convention: each block's @@ -858,9 +898,10 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra build_dspark_markov_head(*this, model, inp_tokens); } + // DFly: re-derive the draft logits through the chained predecessor correction - if (model.dfly_hc_down) { - build_dfly_correction_head(*this, model, inp_tokens); + if (model.dfly_hc_down && getenv("LLAMA_DFLY_CHAIN") != nullptr) { + build_dfly_correction_head(*this, model, inp_tokens, inp_embd_raw); } } diff --git a/tests/gen-tiny-dfly.py b/tests/gen-tiny-dfly.py index e8f2a567a223..44a7b56b20d7 100644 --- a/tests/gen-tiny-dfly.py +++ b/tests/gen-tiny-dfly.py @@ -85,7 +85,7 @@ def t(name: str, shape: tuple[int, ...]) -> None: # DFly encoder: shared base projection + per-draft-layer fusion, then context_norm. # Note there is NO enc.output_norm (hidden_norm): DFly replaces it with context_norm. t("fc.weight", (n_embd, n_feat * n_embd)) - t("layer_fusion.weight", (n_layer, n_feat)) + t("layer_fusion", (n_layer, n_feat)) t("context_norm.weight", (n_embd,)) # TreeFlash predecessor correction (swiglu over [hidden ; prev-token embedding]) From b5055700576f77b90b96594f62be5098c645d54c Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:31:20 -0700 Subject: [PATCH 3/9] speculative: size the draft context round-trip by n_embd_out, not n_embd The fused target context round-trips from the draft encoder's nextn output through a host buffer and back in as an embd batch. Every buffer on that path -- the llama_batch_init embd width, g_embd_buf, verify_g, pending_g_last -- was sized with llama_model_n_embd(model_dft), while llama_decode reads an embd batch on an MTP-type context at llama_model_n_embd_out() per token. Those two are equal for every drafter that emits one context row per token, so this is a no-op for DFlash and DSpark. They are not equal for DFly, which fuses one context per draft layer: n_embd_out is n_layer times wider, so batch.embd was under-allocated by that factor and the decode read past the end of it. The resulting garbage explains both symptoms seen on the reference checkpoint -- draft layers 1..n reading junk context, and a nondeterministic out-of-range index. n_embd_dec only ever describes rows of that nextn buffer (including the DSpark confidence column read through llama_get_embeddings_nextn), so n_embd_out is the right width for all of its uses. --- common/speculative.cpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 1ff0ddeb7e0c..b78cf1f7c5ae 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -430,7 +430,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { // backend sampler chain per seq, attached to ctx_dft std::vector backend_chains; - int32_t n_embd_dec = 0; // draft hidden size + int32_t n_embd_dec = 0; // draft context row width (n_embd_out: per-layer for DFly) int32_t n_embd_enc = 0; // target_layer_ids_n * target_hidden_size int32_t n_embd_tgt = 0; // target model hidden size int32_t n_layer_tgt = 0; // target model layer count @@ -473,7 +473,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { } n_embd_tgt = llama_model_n_embd(model_tgt); - n_embd_dec = llama_model_n_embd(model_dft); + n_embd_dec = llama_model_n_embd_out(model_dft); n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt; n_layer_tgt = llama_model_n_layer(model_tgt); @@ -916,7 +916,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // backend sampler chain per seq, attached to ctx_dft std::vector backend_chains; - int32_t n_embd_dec = 0; // draft hidden size + int32_t n_embd_dec = 0; // draft context row width (n_embd_out: per-layer for DFly) int32_t n_embd_enc = 0; // target_layer_ids_n * target_hidden_size int32_t n_embd_tgt = 0; // target model hidden size @@ -967,7 +967,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } n_embd_tgt = llama_model_n_embd(model_tgt); - n_embd_dec = llama_model_n_embd(model_dft); + n_embd_dec = llama_model_n_embd_out(model_dft); n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt; // read the trained block size from the dflash.block_size metadata key From 26db2347a8d84fb2013bb2de284a0d3645d7dff8 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:50:25 -0700 Subject: [PATCH 4/9] dflash: widen DFly's decoder embd input, and enable the chain by default Root cause of the zero acceptance. A dflash draft context is not MTP-typed -- only COMMON_SPECULATIVE_TYPE_DRAFT_MTP sets that -- so llama_context::decode takes the `hparams.n_embd_inp()` branch when sizing an embd batch, not the `n_embd_out()` one. DFly set only n_embd_out_impl, so the decoder graph built a n_layer*n_embd wide context input while the ubatch supplied n_embd floats per token. Draft layers 1..n-1 read uninitialised memory past the end of it. That one under-allocation produced every symptom: a constant drafted token whose value moved with allocator placement, identical behaviour with the predecessor-correction chain on and off (both were reading the same junk), and the nondeterministic out-of-range index that aborted on Metal. Setting n_embd_inp_impl alongside n_embd_out_impl fixes all of them, and the chain is default-on again -- there is nothing left to gate. Measured on M5 Pro, Metal, public Qwen3-8B bf16 target + AngelSlim Qwen3-8B-DFly-Block8 (revision 5712926) bf16 drafter, block 8, greedy, 64 predicted tokens, llama-server, interleaved A-B-A-B: | arm | tok/s | acceptance | |--------------------|---------------------|-----------------| | AR baseline | 18.35 18.42 18.50 18.52 | -- | | DFly | 38.53 38.65 38.18 38.00 | 52.7% (49/93) | | DFly, chain off | 30.89 | 21.0% (37/176) | 2.08x against a same-path AR baseline whose four readings spread 0.9%. Mean accepted length is 4.50 of a block of 8, which implies an 8-row verify costs about 2.16 plain decode steps -- the right order for a bandwidth-bound bf16 8B target, so the mechanism predicts the size of the win rather than merely accompanying it. The correction chain is worth 21.0% -> 52.7% acceptance and 30.89 -> 38.45 tok/s, so it is load-bearing and its removal costs throughput. Correctness gate: the generated text is byte-identical to the AR baseline at greedy in every speculative arm, which is the guarantee speculative decoding owes and would have caught a mis-indexed block or a corrupted verify. --- src/models/dflash.cpp | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 687ab4b3cae3..c6e63d63dc3e 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -17,12 +17,16 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { hparams.n_embd_inp_enc_impl = (uint32_t) target_layer_ids.size() * hparams.n_embd; // AngelSpec DFly fuses the target context once per DRAFT layer, so the encoder emits - // [n_embd, n_layer] per token instead of a single [n_embd] row. Both the nextn output - // buffer and the embd batch that feeds it back are sized from n_embd_out(), so widening - // it here is the whole host-side contract change. Detected from the fusion tensor rather - // than a KV so a DFly export cannot load as plain DFlash. + // [n_embd, n_layer] per token instead of a single [n_embd] row. Detected from the fusion + // tensor rather than a KV so a DFly export cannot load as plain DFlash. if (ml.get_tensor_meta("layer_fusion")) { + // Both ends of the round-trip must widen: the encoder emits n_layer contexts + // (n_embd_out) and the decoder's embd batch consumes them (n_embd_inp). Setting only + // the former is not enough -- a dflash draft context is not MTP-typed, so + // llama_context::decode sizes an embd batch from n_embd_inp(), and leaving that at + // n_embd hands the graph one layer's context and n_layer-1 layers of junk. hparams.n_embd_out_impl = hparams.n_layer() * hparams.n_embd; + hparams.n_embd_inp_impl = hparams.n_layer() * hparams.n_embd; LLAMA_LOG_INFO("%s: DFly per-layer context fusion (n_layer = %u, n_embd_out = %u)\n", __func__, hparams.n_layer(), hparams.n_embd_out()); } @@ -206,10 +210,6 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { dfly_hc_down = create_tensor(tn(LLM_TENSOR_DFLY_HC_DOWN, "weight"), { n_ff_hc, n_embd }, 0); LLAMA_LOG_INFO("%s: DFly predecessor correction (n_ff = %lld)\n", __func__, (long long) n_ff_hc); - LLAMA_LOG_WARN("%s: DFly predecessor-correction CHAIN IS NOT ENABLED. It is still under " - "development (nondeterministic uninitialised read; see the notes on this branch), " - "so drafts are produced WITHOUT the correction and acceptance will be near zero. " - "Set LLAMA_DFLY_CHAIN=1 to run it anyway.\n", __func__); } else { LLAMA_LOG_INFO("%s: DFly without predecessor correction\n", __func__); } @@ -518,11 +518,10 @@ static void build_dfly_correction_head(llm_graph_context & g, const llama_model return; } - // Only a noise block carries a chain. The draft context also decodes ordinary batches - // (context staging before the K/V injection), and those have no anchor to condition on -- - // their slot 0 is the caller's id_last, which is LLAMA_TOKEN_NULL until the first token is - // committed. Chaining one of those would read token embedding row -1. A block is - // [id_last, MASK, MASK, ...] per sequence, so check that shape before building anything. + // Only a noise block carries a chain. The draft context also decodes ordinary staging + // batches, whose slot 0 is the caller's id_last -- LLAMA_TOKEN_NULL until the first token + // is committed, which would index the embedding table with row -1. A block is + // [id_last, MASK, ...] per sequence, so check that shape before building anything. if (g.ubatch.token == nullptr) { return; } @@ -900,7 +899,9 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra // DFly: re-derive the draft logits through the chained predecessor correction - if (model.dfly_hc_down && getenv("LLAMA_DFLY_CHAIN") != nullptr) { + // LLAMA_DFLY_NO_CHAIN drops the correction, for A/B measurement only: DFly's acceptance + // depends on it, so a run without it is not a meaningful DFly configuration. + if (model.dfly_hc_down && getenv("LLAMA_DFLY_NO_CHAIN") == nullptr) { build_dfly_correction_head(*this, model, inp_tokens, inp_embd_raw); } } From 7e73eebd80b55ab89c27fc4845344e09df354bbe Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Wed, 2 Sep 2026 20:56:29 -0700 Subject: [PATCH 5/9] docs: document DFly and how to tune its draft depth DFly runs through draft-dflash rather than its own --spec-type, which is not obvious from the flag list, so record what it is and how to convert it. Include the revision pin: the reference checkpoint's main states a target depth and vocabulary that do not match its weights, and that loads cleanly before it mis-projects. The tuning note is the part worth having. The predecessor correction runs one full output-head projection per block position, so a draft round costs about fixed + k*per_position while committed tokens saturate with depth. That puts the optimum in the interior, and the default of block_size-1 is past it on the pairing measured here: 42.0 tok/s at n_max 6 against 37.5 at 7, over interleaved rounds. Presented as something to sweep rather than a new default, because the optimum depends on how a backend prices a multi-row verify and so moves with hardware and target. --- docs/speculative.md | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/docs/speculative.md b/docs/speculative.md index 0f9f8a3d977a..bb235ca877cd 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -78,6 +78,45 @@ See: - #22105 +### DFly + +DFly (AngelSpec) is a DFlash variant, so it runs through `draft-dflash` and is detected from the +checkpoint rather than selected with its own `--spec-type`. It differs from plain DFlash in two +ways: the captured target features are fused once per *draft layer* instead of once for the whole +draft model, and a predecessor correction is applied to each block position before the target head, +chained on the token drafted at the previous position. + +Convert it with `--target-model-dir`, as for DFlash. Pin a revision: the reference checkpoint's +`main` states a target depth and vocabulary that do not match its weights, which loads cleanly and +then mis-projects. + +```bash +python convert_hf_to_gguf.py AngelSlim/Qwen3-8B-DFly-Block8 \ + --target-model-dir Qwen/Qwen3-8B --outtype bf16 --outfile Qwen3-8B-DFly.gguf + +llama-server -m Qwen3-8B.gguf -md Qwen3-8B-DFly.gguf \ + --spec-draft-n-max 6 -fa on --jinja +``` + +#### Tuning `--spec-draft-n-max` + +The correction runs one full output-head projection per block position, so a DFly draft round costs +roughly `fixed + k * per_position` while the tokens it commits saturate with depth. The optimum is +therefore interior, and the default (the trained block size minus one) is not always it. + +Measured on an M5 Pro with the pairing above, greedy, interleaved rounds: + +| `--spec-draft-n-max` | tok/s | acceptance | committed tokens per round | +|---|---|---|---| +| 5 | 35.7 | 57.5% | 4.00 | +| 6 | 42.0 | 64.9% | 4.99 | +| 7 (default for block size 8) | 37.5 | 52.7% | 4.82 | + +Sweep it rather than assuming the largest value wins. The optimum depends on how the backend prices +a multi-row verify, so it moves with hardware and with the target, and is not a property of the +drafter alone. Note also that changing the value changes the shape of the drafted block, so the +drafts differ entirely between settings rather than simply being truncated. + ### DSpark (`draft-dspark`) DSpark extends DFlash with a semi-autoregressive _Markov head_: the draft still emits a whole From b995a032e0f839681ca1f9815ef6f54f22270026 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Thu, 3 Sep 2026 08:27:53 -0700 Subject: [PATCH 6/9] dflash: name the correction-without-fusion lineage instead of failing on a count A drafter can carry a trained predecessor correction without the per-draft-layer fusion: DSpark plus correction, tap_fusion none. That is a real checkpoint shape, not a hypothetical one. DFly-ness is detected from layer_fusion, and the hidden_correction tensors are only created on that path, so such a checkpoint loads as plain DSpark, leaves its correction weights uncreated, and fails in done_getting_tensors with "wrong number of tensors; expected N, got M". That is the same opaque failure that already cost time once on this arch, where a 79-versus-75 count turned out to be the log-SNR pair rather than the markov head it looked like. Reject it at load with the reason instead. Supporting the lineage properly would mean deciding how a markov bias and a predecessor correction compose on the same block, which is not something to guess at, so this asserts the boundary rather than inventing behaviour behind it. Verified on a purpose-built tiny GGUF (hidden_correction present, layer_fusion absent): the named error fires where the tensor-count mismatch used to. Valid DFly still loads, the fusion parity test still passes, and the pre-existing fusion-plus-markov conflict still reports as mutually exclusive. --- src/models/dflash.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index c6e63d63dc3e..87d1d567db8d 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -139,6 +139,16 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { "The reference chain runs the full target head"); } + // A predecessor correction without the per-layer fusion is a third lineage (DSpark plus + // correction, tap_fusion none). It is not supported here, and the correction weights are + // only created on the DFly path, so letting it through means an opaque + // "wrong number of tensors" from done_getting_tensors rather than a reason. + if (!is_dfly && ml->get_tensor_meta("hidden_correction.down.weight")) { + throw std::runtime_error("dflash: checkpoint carries hidden_correction weights but no layer_fusion. " + "That lineage (DSpark plus predecessor correction) is not supported; " + "loading it as plain DSpark would drop the trained correction"); + } + // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head // // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) From 63c9d09237547da9814329e527ff38b3ae8be127 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Fri, 4 Sep 2026 00:08:49 -0700 Subject: [PATCH 7/9] dflash: validate DFly target_* against the target, and drop disabled correction weights Two review points on the DFly converter. The target_layer_ids bound could not catch the case it was written for. The reference checkpoint's main-branch config declares depth 80 and vocab 120832 against a real 36 and 151936, and capture ids up to 33 sit inside the declared 80, so the check passed and the drafter mis-projected. A declared depth cannot police itself. The declared target_* are now compared against the target model read from --target-model-dir, which is the only authoritative shape here, and the ids are bounded by the real depth. Shapes are read per key, so one unusable value does not disable the other comparisons, and a target config that is missing, unparseable or not a JSON object leaves the comparison inert rather than raising here; set_vocab already reads the same file and fails on it. hidden_correction.* weights were converted even when the config disables correction. The runtime turns the feature on from the presence of hidden_correction.down.weight, so a checkpoint carrying stale weights silently re-enabled a feature its config had switched off. They are now omitted when enable_hidden_correction is false, with one log line saying so. There is no metadata key for this, tensor presence is the only signal, so omitting the tensors is the whole fix. Verified by running the real classes over synthetic configs: the 80/120832 config is accepted before this change and rejected after, a truthful config still converts, ids past the real depth are caught with the declared depth absent, a null value in the target config no longer hides the other lies, a nested text_config is followed, and a non-object or corrupt target config is handled without a traceback. Omitting --target-model-dir preserves the old behaviour. On the correction path the enabled case and unrelated tensors are identical with and without the change; only the disabled case differs. flake8 clean, and ty reports the same 8 pre-existing diagnostics before and after. --- conversion/qwen.py | 46 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/conversion/qwen.py b/conversion/qwen.py index dde22ca348cd..a553131fd764 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -757,6 +757,24 @@ def __init__(self, dir_model, *args, **kwargs): # clean and then mis-projects. Refuse instead of converting a checkpoint that lies. target_layers = hp.get("target_num_hidden_layers") layer_ids = hp.get("target_layer_ids") or [] + + # A declared depth cannot police itself: the reference config's 80 is self-consistent with + # capture ids up to 33 and still wrong. The target model is the only authoritative shape, + # so compare the declared target_* against it and bound the ids by the real depth. + real = self._target_shapes() + for cfg_key, hp_key in (("num_hidden_layers", "target_num_hidden_layers"), + ("vocab_size", "target_vocab_size"), + ("hidden_size", "target_hidden_size")): + declared = hp.get(hp_key) + if declared is not None and cfg_key in real and int(declared) != real[cfg_key]: + raise ValueError( + f"DFly {hp_key} is {declared} but the target model reports " + f"{cfg_key}={real[cfg_key]}. The config metadata does not describe the target " + "-- pin a known-good revision (the reference checkpoint's is 5712926)." + ) + if "num_hidden_layers" in real: + target_layers = real["num_hidden_layers"] + if target_layers and layer_ids and max(layer_ids) >= int(target_layers): raise ValueError( f"DFly target_layer_ids {layer_ids} exceed target_num_hidden_layers {target_layers}. " @@ -785,6 +803,26 @@ def __init__(self, dir_model, *args, **kwargs): # slot 0 of a DFly block is the committed bonus anchor, not a prediction slot self._sample_from_anchor = not bool(hp.get("dspark_bonus_anchor", True)) + def _target_shapes(self) -> dict[str, int]: + """Shapes read from --target-model-dir, for the keys it declares. Empty when unavailable.""" + if self.target_model_dir is None: + return {} # set_vocab raises on this later; nothing authoritative to compare against + try: + with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f: + cfg = json.load(f) + except (OSError, ValueError): + return {} # unreadable; set_vocab reads the same file and raises + if not isinstance(cfg, dict): + return {} + cfg = {**cfg, **(cfg.get("text_config") or {})} + shapes = {} + for k in ("num_hidden_layers", "vocab_size", "hidden_size"): + try: + shapes[k] = int(cfg[k]) + except (KeyError, ValueError, TypeError): + continue # one unusable value must not disable the other comparisons + return shapes + def set_gguf_parameters(self): super().set_gguf_parameters() self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor) @@ -798,9 +836,17 @@ def prepare_tensors(self): ) _seen_correction = False + _dropped_correction = False def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: if "hidden_correction." in name: + # the runtime turns correction on from the presence of hidden_correction.down.weight, + # so shipping these while the config disables it silently re-enables the feature + if not self._has_correction: + if not self._dropped_correction: + logger.info("DFly: enable_hidden_correction is false, dropping hidden_correction.* weights") + self._dropped_correction = True + return self._seen_correction = True yield from super().modify_tensors(data_torch, name, bid) From 844b8ce746a0fc6da87339a88ba94262029fb9d5 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:29:37 -0700 Subject: [PATCH 8/9] dflash: check the DFly hidden size against the target, not against itself An omitted target_hidden_size defaulted to the drafter's own hidden_size, so the residual-fusion check compared a value with itself and passed. The target model's config is already read for the other target_* comparisons; use its hidden_size when it is available. Also trims the per-layer fusion comment to the invariant. --- conversion/qwen.py | 9 ++++++--- src/models/dflash.cpp | 7 ++----- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/conversion/qwen.py b/conversion/qwen.py index a553131fd764..d4d502965d6a 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -782,10 +782,13 @@ def __init__(self, dir_model, *args, **kwargs): "(the reference checkpoint's is 5712926)." ) - if int(hp.get("target_hidden_size", hp["hidden_size"])) != int(hp["hidden_size"]): + # an omitted target_hidden_size would default to the draft width and compare equal to itself, + # so prefer the target model's own hidden_size whenever it is readable + target_hidden = real.get("hidden_size", hp.get("target_hidden_size", hp["hidden_size"])) + if int(target_hidden) != int(hp["hidden_size"]): raise ValueError( - "DFly residual fusion requires target_hidden_size == hidden_size, got " - f"{hp.get('target_hidden_size')} vs {hp['hidden_size']}." + "DFly residual fusion requires the target hidden_size to equal the drafter's, got " + f"{target_hidden} vs {hp['hidden_size']}." ) if hp.get("markov_rank") or hp.get("enable_confidence_head"): diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 87d1d567db8d..155e119f71a0 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -20,11 +20,8 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { // [n_embd, n_layer] per token instead of a single [n_embd] row. Detected from the fusion // tensor rather than a KV so a DFly export cannot load as plain DFlash. if (ml.get_tensor_meta("layer_fusion")) { - // Both ends of the round-trip must widen: the encoder emits n_layer contexts - // (n_embd_out) and the decoder's embd batch consumes them (n_embd_inp). Setting only - // the former is not enough -- a dflash draft context is not MTP-typed, so - // llama_context::decode sizes an embd batch from n_embd_inp(), and leaving that at - // n_embd hands the graph one layer's context and n_layer-1 layers of junk. + // both ends must widen: a dflash draft context is not MTP-typed, so it sizes its embd batch + // from n_embd_inp, and setting only n_embd_out gives the graph one layer's context plus junk hparams.n_embd_out_impl = hparams.n_layer() * hparams.n_embd; hparams.n_embd_inp_impl = hparams.n_layer() * hparams.n_embd; LLAMA_LOG_INFO("%s: DFly per-layer context fusion (n_layer = %u, n_embd_out = %u)\n", From 1ecc8d65d9ce6878a5b090a64f35795f1245ce0e Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:01:44 -0700 Subject: [PATCH 9/9] dflash: reject unsupported DFly anchor convention --- conversion/qwen.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/conversion/qwen.py b/conversion/qwen.py index d4d502965d6a..9876e567453a 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -803,8 +803,9 @@ def __init__(self, dir_model, *args, **kwargs): if correction_type != "swiglu": raise ValueError(f"unsupported hidden_correction_type {correction_type!r} (only 'swiglu')") - # slot 0 of a DFly block is the committed bonus anchor, not a prediction slot - self._sample_from_anchor = not bool(hp.get("dspark_bonus_anchor", True)) + if not bool(hp.get("dspark_bonus_anchor", True)): + raise ValueError("DFly requires dspark_bonus_anchor=true; slot 0 is the committed anchor") + self._sample_from_anchor = False def _target_shapes(self) -> dict[str, int]: """Shapes read from --target-model-dir, for the keys it declares. Empty when unavailable."""