dflash: AngelSpec DFly drafter support by bri-prism · Pull Request #146 · PrismML-Eng/llama.cpp

bri-prism · 2026-09-03T02:18:49Z

What

Adds AngelSpec DFly drafter support to the dflash arch, so a DFly checkpoint can be
converted and used as a speculative drafter. Verified against the public reference pairing:
Qwen3-8B as target and AngelSlim/Qwen3-8B-DFly-Block8 as drafter.

Also fixes one bug in common/speculative.cpp that is independent of DFly and could be
taken on its own (second commit).

Why

DFly extends the DFlash draft path with two pieces we did not have:

Per-draft-layer context fusion. DFlash projects the captured target features once and
injects that single context into every draft layer. DFly keeps that shared projection as a
base context, then adds a per-draft-layer softmax-weighted mix of the raw per-target-layer
features, so each draft layer gets its own context.
A TreeFlash predecessor correction applied before the target head, 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.

The expensive prerequisite was already here. dflash already reads a target_layers array
and sizes its encoder input as |target_layers| * n_embd, which is exactly the concatenated
feature layout DFly consumes, so multi-layer capture needed no work.

How

DFly is detected from the presence of the layer_fusion tensor rather than from a metadata
key, so a DFly export cannot quietly load as plain DFlash. A checkpoint that declares both a
DFly fusion and a DSpark Markov head is rejected with that reason instead of failing later on
a missing tensor. Reduced draft vocabularies are rejected, since the reference chain runs the
full target head.

Nothing was needed in the draft reader. DFly reports no Markov head, so it already takes the
DFlash path (block rows 1 through n-1) and the plain per-token probability cut.

Two width contracts have to widen together for the fused context to survive its round trip
through the host, and this is the part worth reviewing:

n_embd_out sizes the encoder's nextn output. common/speculative.cpp was sizing every
buffer on that path with llama_model_n_embd(model_dft) instead. Those two are equal for
any drafter that emits one context row per token, so that commit is a no-op for DFlash and
DSpark, but they are not equal for DFly and the batch was under-allocated.
n_embd_inp sizes the embd batch that feeds the context back in. A dflash draft context is
not MTP-typed, so llama_context::decode takes the n_embd_inp() branch. Setting only
n_embd_out leaves the decoder reading one layer's context and n_layer-1 layers of junk.

Measurements

M5 Pro, Metal, bf16 target and bf16 drafter, block 8, greedy, 64 predicted tokens, through
llama-server, with the arms interleaved A-B-A-B:

arm	tok/s	acceptance
autoregressive baseline	18.35, 18.42, 18.50, 18.52	n/a
DFly	38.53, 38.65, 38.18, 38.00	52.7 percent (49/93)
DFly, correction disabled	30.89	21.0 percent (37/176)

That is 2.08x against a same-path baseline whose four readings spread 0.9 percent. Mean
accepted length is 4.50 of a block of 8, which puts an 8-row verify at about 2.16 plain decode
steps. That is the right order for a bandwidth-bound bf16 8B target, so the mechanism accounts
for the size of the speedup rather than just accompanying it.

The correction is load-bearing: disabling it drops acceptance from 52.7 to 21.0 percent.
LLAMA_DFLY_NO_CHAIN=1 exists to make that A/B measurable and is not a normal setting.

Testing

Generated text is byte identical to the autoregressive baseline at greedy in every
speculative arm. That is the guarantee speculative decoding owes, and it is what would catch
a mis-indexed block or a corrupted verify.
tests/test-dfly-fusion.cpp pins the fusion math and the resulting memory layout against an
independent scalar reference. The axis handling needs two reshape and permute round trips,
and a transposed axis would otherwise still load and still draft plausibly. The test also
asserts the per-layer contexts actually differ, so a collapsed broadcast fails.
tests/gen-tiny-dfly.py builds a tiny random DFly GGUF for loader smoke tests, including
the rejection paths above.

Notes for review

The conversion path covers Qwen3DFlyModel and the Qwen3DSparkDFlareV2Model alias, and
refuses a config whose target_layer_ids exceed its stated target depth. The published
metadata of the reference checkpoint does exactly that on its main branch, and it loads clean
and then mis-projects, so pin a known good revision when converting.

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.

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.

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.

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.

Copilot

🟡 Changes recommended

Target validation, disabled correction handling, and predecessor sampling are currently inconsistent with the intended behavior.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds AngelSpec DFly speculative drafter support to the DFlash architecture.

Changes:

Implements per-layer context fusion and predecessor correction.
Adds conversion and GGUF tensor mappings.
Fixes speculative context-width allocation and adds fusion tests.

File summaries

File	Description
`tests/test-dfly-fusion.cpp`	Tests fusion math and layout.
`tests/gen-tiny-dfly.py`	Generates a tiny DFly GGUF.
`tests/CMakeLists.txt`	Registers the fusion test.
`src/models/dflash.cpp`	Implements DFly loading and graphs.
`src/llama-model.h`	Adds DFly tensor fields.
`src/llama-arch.h`	Declares DFly tensor identifiers.
`src/llama-arch.cpp`	Registers DFly tensor metadata.
`gguf-py/gguf/tensor_mapping.py`	Maps source DFly tensor names.
`gguf-py/gguf/constants.py`	Adds DFly GGUF constants.
`conversion/qwen.py`	Adds DFly conversion support.
`conversion/__init__.py`	Registers DFly model classes.
`common/speculative.cpp`	Uses the encoder output width for buffers.

Review details

Files reviewed: 12/12 changed files
Comments generated: 3
Review effort level: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+        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):


+        if "hidden_correction." in name:
+            self._seen_correction = True
+
+        yield from super().modify_tensors(data_torch, name, bid)


+        // greedy chain: the next position conditions on this position's drafted token
+        if (i + 1 < block_drafts) {
+            prev = ggml_argmax(ctx0, col);
+        }


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.

… 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.

…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.

bri-prism added 4 commits September 2, 2026 16:19

github-actions Bot added testing model conversion labels Sep 3, 2026

bri-prism requested a review from khosravipasha September 3, 2026 02:48

khosravipasha requested a balanced review from Copilot September 3, 2026 03:29

Copilot started reviewing on behalf of khosravipasha September 3, 2026 03:30 View session

Copilot AI reviewed Sep 3, 2026

View reviewed changes

github-actions Bot added the documentation Improvements or additions to documentation label Sep 3, 2026

bri-prism added 2 commits September 3, 2026 08:27

Navigation Menu

Uh oh!

dflash: AngelSpec DFly drafter support - #146

dflash: AngelSpec DFly drafter support#146
bri-prism wants to merge 7 commits into
prism-v7from
feat/dfly-drafter

bri-prism commented Sep 3, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

bri-prism commented Sep 3, 2026

What

Why

How

Measurements

Testing

Notes for review

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants