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 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/conversion/qwen.py b/conversion/qwen.py index ff72c61314f0..a553131fd764 100644 --- a/conversion/qwen.py +++ b/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", 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 diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 7777e7efa94d..f05f1b740ee4 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..87d1d567db8d 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -16,6 +16,21 @@ 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. 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()); + } + // 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 +129,26 @@ 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") != 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"); + } + + // 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) @@ -129,6 +164,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 +197,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), { 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 +335,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 +490,138 @@ 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_tensor * inp_embd_raw) { + 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; + } + + // 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; + } + + 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; + 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 hidden_stride = (size_t) block_drafts * hidden->nb[1]; + const size_t base_stride = (size_t) block_drafts * base->nb[1]; + + // 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 + 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))] + 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])); + + 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 @@ -413,11 +646,16 @@ 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) { - 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 +666,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); @@ -499,6 +746,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 @@ -653,6 +906,14 @@ 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 + // 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); + } } // DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above): diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index a53b5c693f8b..b90422641660 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..44a7b56b20d7 --- /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", (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; +}