dflash: AngelSpec DFly drafter support

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<llama_sampler *> 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<llama_sampler *> 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

conversion/__init__.py

@@ -54,6 +54,8 @@
     "DeepseekV3ForCausalLM": "deepseek",
     "DeepseekV32ForCausalLM": "deepseek",
     "DFlashDraftModel": "qwen",
+    "Qwen3DFlyModel": "qwen",
+    "Qwen3DSparkDFlareV2Model": "qwen",
     "Qwen3DSparkModel": "qwen",
     "DSparkDraftModel": "qwen",
     "DSparkSpeculator": "qwen",

conversion/qwen.py

@@ -729,6 +729,129 @@ 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 []
+
+        # 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}. "
+                "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 _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)
+
+    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
+    _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)
+
+
 @ModelBase.register(
     "Qwen3DSparkModel",
     "DSparkDraftModel",

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

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,
     ],

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: (

src/llama-arch.cpp

@@ -657,6 +657,13 @@ static const std::map<llm_tensor, const char *> 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, llm_tensor_info> 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}},

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,
 };
 
 

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<int32_t> target_layer_ids;