From f4e67e05d279974214d6340c3cbe451d0c1699f8 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Wed, 8 Jul 2026 16:23:15 +0800 Subject: [PATCH 01/10] dspark : support Gemma4 backbone drafts (WIP) Assisted-by: Claude Fable 5 --- conversion/__init__.py | 1 + conversion/gemma.py | 65 +++++++++++++++++++++++++++++++++++++++ gguf-py/gguf/constants.py | 5 +++ src/models/dflash.cpp | 62 ++++++++++++++++++++++++++++++++----- 4 files changed, 125 insertions(+), 8 deletions(-) diff --git a/conversion/__init__.py b/conversion/__init__.py index 3232a1050bfe..594151e8ab05 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -88,6 +88,7 @@ "Gemma3nForCausalLM": "gemma", "Gemma3nForConditionalGeneration": "gemma", "Gemma4AssistantForCausalLM": "gemma", + "Gemma4DSparkModel": "gemma", "Gemma4ForConditionalGeneration": "gemma", "Gemma4ForCausalLM": "gemma", "Gemma4UnifiedForConditionalGeneration": "gemma", diff --git a/conversion/gemma.py b/conversion/gemma.py index 6b4d7d17154d..ffe1dff282a7 100644 --- a/conversion/gemma.py +++ b/conversion/gemma.py @@ -851,6 +851,71 @@ def set_gguf_parameters(self): self.gguf_writer.add_nextn_predict_layers(self.block_count) +@ModelBase.register("Gemma4DSparkModel") +class Gemma4DSparkModel(Gemma4Model): + # DSpark draft with a Gemma4 backbone; same DeepSpec flat config schema and shared + model_arch = gguf.MODEL_ARCH.DFLASH + + def set_vocab(self): + if self.target_model_dir is None: + raise ValueError( + "DSpark draft model requires --target-model-dir to be specified. " + "Please provide the path to the target model directory containing the tokenizer." + ) + logger.info(f"DSpark: Using tokenizer from target model: {self.target_model_dir}") + original_dir = self.dir_model + self.dir_model = self.target_model_dir + super().set_vocab() + self.dir_model = original_dir + + mask_token_id = self.hparams.get("mask_token_id") + if mask_token_id is not None: + self.gguf_writer.add_mask_token_id(mask_token_id) + + def set_gguf_parameters(self): + assert not self.hparams.get("enable_moe_block", False), "Gemma4 DSpark with MoE blocks is not supported" + if self.hparams.get("markov_rank", 0) > 0: + assert self.hparams.get("markov_head_type") == "vanilla", "only the vanilla Markov head is supported" + + # the draft uses full attention on every layer + assert all(lt == "full_attention" for lt in self.hparams["layer_types"]) + self.hparams["sliding_window_pattern"] = 1 + + super().set_gguf_parameters() + + self.gguf_writer.add_block_size(self.hparams.get("block_size", 7)) + + # flat DeepSpec schema; mirror DFlash's +1 extract-layer convention + target_layer_ids = self.hparams.get("target_layer_ids", []) + if target_layer_ids: + extract_layer_ids = [i + 1 for i in target_layer_ids] + self.gguf_writer.add_target_layers(extract_layer_ids) + + self.gguf_writer.add_markov_rank(self.hparams.get("markov_rank", 0)) + + # Gemma4TextScaledWordEmbedding scales the shared token embeddings by sqrt(hidden_size) + self.gguf_writer.add_embedding_scale(self.hparams["hidden_size"] ** 0.5) + # Gemma4DSparkAttention uses scaling = 1.0 + self.gguf_writer.add_attention_scale(1.0) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + # embed_tokens / lm_head are byte-identical to the target and shared at runtime -- drop them + if name.endswith(("embed_tokens.weight", "lm_head.weight")): + return None + if not name.startswith("model."): + name = "model." + name + return super().filter_tensors((name, gen)) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # "post_attention_layernorm" is ambiguous in the shared DFlash tensor map -- resolve it explicitly for the Gemma4 backbone + if name.endswith(".post_attention_layernorm.weight"): + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_POST_NORM, bid), data_torch) + return + yield from super().modify_tensors(data_torch, name, bid) + + @ModelBase.register("Gemma4ForConditionalGeneration") @ModelBase.example("google/gemma-4-31B-it", "google/gemma-4-26B-A4B-it", "google/gemma-4-E2B-it") class Gemma4VisionAudioModel(MmprojModel): diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d043c9b6ecce..a81b0b6c142a 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -4822,6 +4822,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, @@ -4829,6 +4830,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.ATTN_OUT, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_POST_NORM, MODEL_TENSOR.ATTN_SINKS, MODEL_TENSOR.ATTN_Q_A, MODEL_TENSOR.ATTN_Q_B, @@ -4847,9 +4849,11 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.HC_HEAD_BASE, MODEL_TENSOR.HC_HEAD_SCALE, MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_PRE_NORM, MODEL_TENSOR.FFN_GATE, MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_POST_NORM, MODEL_TENSOR.FFN_GATE_INP, MODEL_TENSOR.FFN_EXP_PROBS_B, MODEL_TENSOR.FFN_GATE_EXP, @@ -4858,6 +4862,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_GATE_SHEXP, MODEL_TENSOR.FFN_DOWN_SHEXP, MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.LAYER_OUT_SCALE, MODEL_TENSOR.FC, MODEL_TENSOR.ENC_OUTPUT_NORM, MODEL_TENSOR.D2T, diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 5b70a5179496..5fd961b5f5b4 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -76,6 +76,12 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; } + // Gemma4 backbone: scaled token embeddings, kq_scale override and final-logit softcapping + hparams.f_final_logit_softcapping = 0.0f; + ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false); + ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false); + ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false); + type = LLM_TYPE_UNKNOWN; } @@ -177,7 +183,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head_k * n_head }, 0); layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), { n_embd, n_embd_k_gqa }, 0); - layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, 0); + layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, TENSOR_NOT_REQUIRED); // absent for Gemma4 (V = K projection) layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0); layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0); @@ -187,6 +193,20 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0); layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0); layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0); + + // Gemma4 backbone: sandwich norms, per-layer output scalar and proportional-rope freq factors + layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED); + layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED); + layer.out_scale = create_tensor(tn(LLM_TENSOR_LAYER_OUT_SCALE, "weight", i), { 1u }, TENSOR_NOT_REQUIRED); + layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), { n_embd_head_k/2 }, + TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); + + // the optional tensors must agree on the backbone type: + // Gemma4 has sandwich norms and no V projection (V = K), Qwen has the opposite + const bool gemma = layer.attn_post_norm != nullptr; + if (gemma == (layer.wv != nullptr) || gemma != (layer.ffn_post_norm != nullptr)) { + throw std::runtime_error(format("inconsistent DFlash backbone tensors in layer %d", i)); + } } } @@ -367,7 +387,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra inp_attn = build_attn_inp_kv(); } - const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); + const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; // KV cache injection if (ubatch.embd) { @@ -385,14 +405,18 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra 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); + ggml_tensor * Vcur = layer.wv ? build_lora_mm(layer.wv, inp_g) : Kcur; 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); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + if (layer.wv == nullptr) { + // Gemma4: V is the K projection passed through a scale-less RMS norm, no rope + Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); + } Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, + ctx0, Kcur, inp_pos, layer.rope_freqs, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow ); @@ -453,6 +477,9 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra ggml_tensor * inp_tokens = inp->tokens; ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens); + if (hparams.f_embedding_scale != 0.0f) { + inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale); + } cb(inpL, "inp_noise_embd", -1); res->add_input(std::move(inp)); @@ -465,7 +492,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm); ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm); - ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm); + ggml_tensor * Vcur = layer.wv ? build_lora_mm(layer.wv, noise_norm) : Kcur; Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); @@ -473,14 +500,18 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + if (layer.wv == nullptr) { + // Gemma4: V is the K projection passed through a scale-less RMS norm, no rope + Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); + } Qcur = ggml_rope_ext( - ctx0, Qcur, inp_pos, nullptr, + ctx0, Qcur, inp_pos, layer.rope_freqs, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow ); Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, + ctx0, Kcur, inp_pos, layer.rope_freqs, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow ); @@ -493,21 +524,36 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il) : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + if (layer.attn_post_norm) { + cur = build_norm(cur, layer.attn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_post_norm", il); + } + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); cb(ffn_inp, "ffn_inp", il); cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); + // sandwich norms mark a Gemma4 backbone, which uses GELU instead of SiLU cur = build_ffn(cur, layer.ffn_up, NULL, layer.ffn_up_s, layer.ffn_gate, NULL, layer.ffn_gate_s, layer.ffn_down, NULL, layer.ffn_down_s, NULL, - LLM_FFN_SILU, LLM_FFN_PAR, il); + layer.ffn_post_norm ? LLM_FFN_GELU : LLM_FFN_SILU, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); + if (layer.ffn_post_norm) { + cur = build_norm(cur, layer.ffn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_post_norm", il); + } + cur = ggml_add(ctx0, cur, ffn_inp); + + if (layer.out_scale) { + cur = ggml_mul(ctx0, cur, layer.out_scale); + } cb(cur, "l_out", il); inpL = cur; From 1849dd08727b4cf444434d0a4232a739254bfa94 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Sun, 12 Jul 2026 09:45:11 +0800 Subject: [PATCH 02/10] Fix GGUF conversion error --- conversion/gemma.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/conversion/gemma.py b/conversion/gemma.py index ffe1dff282a7..fa454417e22c 100644 --- a/conversion/gemma.py +++ b/conversion/gemma.py @@ -891,8 +891,6 @@ def set_gguf_parameters(self): extract_layer_ids = [i + 1 for i in target_layer_ids] self.gguf_writer.add_target_layers(extract_layer_ids) - self.gguf_writer.add_markov_rank(self.hparams.get("markov_rank", 0)) - # Gemma4TextScaledWordEmbedding scales the shared token embeddings by sqrt(hidden_size) self.gguf_writer.add_embedding_scale(self.hparams["hidden_size"] ** 0.5) # Gemma4DSparkAttention uses scaling = 1.0 From 3e8108f3f4a9d3cbc0019df55cd4fdd2ec11f4e4 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Wed, 29 Jul 2026 13:46:00 +0800 Subject: [PATCH 03/10] dflash: hoist the backbone FFN activation choice out of the layer loop Co-authored-by: Zaire404 <3147879462@qq.com> Assisted-by: Claude --- src/models/dflash.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 5fd961b5f5b4..89aec57e555e 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -389,6 +389,9 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + // the Gemma4 backbone (marked by its sandwich norms) uses GELU instead of SiLU + const llm_ffn_op_type ffn_act = model.layers[0].ffn_post_norm ? LLM_FFN_GELU : LLM_FFN_SILU; + // KV cache injection if (ubatch.embd) { auto inp = std::make_unique(n_embd); @@ -535,13 +538,12 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); - // sandwich norms mark a Gemma4 backbone, which uses GELU instead of SiLU cur = build_ffn(cur, layer.ffn_up, NULL, layer.ffn_up_s, layer.ffn_gate, NULL, layer.ffn_gate_s, layer.ffn_down, NULL, layer.ffn_down_s, NULL, - layer.ffn_post_norm ? LLM_FFN_GELU : LLM_FFN_SILU, LLM_FFN_PAR, il); + ffn_act, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); if (layer.ffn_post_norm) { From 9ba014181e753da06bfde5976305b1b6c170a193 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Wed, 29 Jul 2026 13:59:06 +0800 Subject: [PATCH 04/10] dspark: update the docs for the p-min fold and Gemma4 backbone, drop a stale TODO Co-authored-by: Zaire404 <3147879462@qq.com> Assisted-by: Claude --- docs/speculative.md | 8 ++++---- src/models/dflash.cpp | 3 --- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/docs/speculative.md b/docs/speculative.md index 0f9f8a3d977a..257256524225 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -100,11 +100,11 @@ llama-server -m Qwen3-4B.gguf -md Qwen3-4B-DSpark.gguf \ `--spec-draft-n-max` is clamped to the draft model's trained block size. -`--spec-draft-conf-min P` truncates each drafted block at the first position whose predicted -acceptance (from the draft's confidence head, if present) falls below `P` (default 0 = disabled). +`--spec-draft-p-min P` also gates the DSpark confidence head: each drafted block is truncated at +the first position whose predicted acceptance falls below `P` (default 0 = disabled). -Currently only drafts with a Qwen3 backbone are supported; support for other backbones -(e.g. Gemma4) is planned. +DSpark drafts come with either a Qwen3-style or a Gemma4-style backbone; the backbone is detected +from the checkpoint, and both convert and run the same way. DSpark drafts exported in the [speculators](https://github.com/vllm-project/speculators) format (for example [`RedHatAI/gemma-4-31B-it-speculator.dspark`](https://huggingface.co/RedHatAI/gemma-4-31B-it-speculator.dspark)) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 89aec57e555e..c41f72bbb33a 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -102,9 +102,6 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { } // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head - // - // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) - // need their own conversion path and graph tweaks const struct ggml_tensor * markov_meta = ml->get_tensor_meta("markov_w1.weight"); if (markov_meta) { const int64_t dspark_markov_rank = markov_meta->ne[0]; From 75cf5ef1a2e0566a84f3129280f46ff8dc7a4fbc Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Wed, 5 Aug 2026 15:33:02 +0800 Subject: [PATCH 05/10] dspark: split the Gemma4 draft graph into a separate class Assisted-by: Claude Fable --- src/models/dflash.cpp | 227 ++++++++++++++++++++++++++++++++++++++---- src/models/models.h | 4 + 2 files changed, 209 insertions(+), 22 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index c41f72bbb33a..a2a732a6364d 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -216,6 +216,9 @@ std::unique_ptr llama_model_dflash::build_arch_graph(const ll if (hparams.dsv4_hc_mult > 0) { return std::make_unique(*this, params); } + if (layers[0].attn_post_norm) { + return std::make_unique(*this, params); + } return std::make_unique>(*this, params); default: GGML_ABORT("invalid graph type"); @@ -384,10 +387,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra inp_attn = build_attn_inp_kv(); } - const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; - - // the Gemma4 backbone (marked by its sandwich norms) uses GELU instead of SiLU - const llm_ffn_op_type ffn_act = model.layers[0].ffn_post_norm ? LLM_FFN_GELU : LLM_FFN_SILU; + const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); // KV cache injection if (ubatch.embd) { @@ -405,16 +405,206 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra const auto & layer = model.layers[il]; ggml_tensor * Kcur = build_lora_mm(layer.wk, inp_g); - ggml_tensor * Vcur = layer.wv ? build_lora_mm(layer.wv, inp_g) : Kcur; + ggml_tensor * Vcur = build_lora_mm(layer.wv, inp_g); 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); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - if (layer.wv == nullptr) { - // Gemma4: V is the K projection passed through a scale-less RMS norm, no rope - Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + cb(Kcur, "Kcur_injected", il); + cb(Vcur, "Vcur_injected", il); + + if (use_iswa) { + // route each layer's K/V to its sub-cache: SWA layers -> sliding cache, full -> dense + const bool is_swa = hparams.is_swa(il); + const auto * kv = is_swa ? inp_attn_iswa->mctx->get_swa() : inp_attn_iswa->mctx->get_base(); + ggml_tensor * k_idxs = is_swa ? inp_attn_iswa->get_k_idxs_swa() : inp_attn_iswa->get_k_idxs(); + ggml_tensor * v_idxs = is_swa ? inp_attn_iswa->get_v_idxs_swa() : inp_attn_iswa->get_v_idxs(); + // rotate K/V into the cache's rotated space + ggml_tensor * k_rot = is_swa ? inp_attn_iswa->self_k_rot_swa : inp_attn_iswa->self_k_rot; + ggml_tensor * v_rot = is_swa ? inp_attn_iswa->self_v_rot_swa : inp_attn_iswa->self_v_rot; + if (k_rot) { + Kcur = llama_mul_mat_hadamard(ctx0, Kcur, k_rot); + } + if (v_rot) { + Vcur = llama_mul_mat_hadamard(ctx0, Vcur, v_rot); + } + ggml_build_forward_expand(gf, kv->cpy_k(ctx0, Kcur, k_idxs, il)); + ggml_build_forward_expand(gf, kv->cpy_v(ctx0, Vcur, v_idxs, il)); + } else { + // rotate K/V into the cache's rotated space + if (inp_attn->self_k_rot) { + Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot); + } + if (inp_attn->self_v_rot) { + Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot); + } + ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il)); + ggml_build_forward_expand(gf, inp_attn->mctx->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il)); } + } + + res->t_embd = inp_g; + + ggml_build_forward_expand(gf, inp_g); + return; + } + + // tok_embd from the target model (shared via ctx_other) + auto * tok_embd = model.tok_embd; + if (tok_embd == nullptr) { + GGML_ASSERT(cparams.ctx_other != nullptr); + const auto * model_other = llama_get_model(cparams.ctx_other); + + GGML_ASSERT(model_other->tok_embd != nullptr && "DFlash decoder requires the target model's token embeddings"); + tok_embd = model_other->tok_embd; + } + + auto inp = std::make_unique(n_embd); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + ggml_tensor * inp_tokens = inp->tokens; + + ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens); + cb(inpL, "inp_noise_embd", -1); + + res->add_input(std::move(inp)); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + ggml_tensor * noise_norm = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(noise_norm, "noise_norm", il); + + ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm); + ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm); + ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm); + + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); + 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); + + Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); + Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + // cache-aware, non-causal attention + ggml_tensor * cur = use_iswa + ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il) + : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + layer.ffn_up, NULL, NULL, + layer.ffn_gate, NULL, NULL, + layer.ffn_down, NULL, NULL, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "l_out", il); + + inpL = cur; + } + + ggml_tensor * cur = build_norm(inpL, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + + res->t_embd = cur; + + // lm_head from the target model (shared via ctx_other) + auto * output = model.output; + if (output == nullptr) { + GGML_ASSERT(cparams.ctx_other != nullptr); + const auto * model_other = llama_get_model(cparams.ctx_other); + GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection"); + output = model_other->output; + } + + cur = build_lora_mm(output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); + + // DSpark: bias the draft logits with the Markov head + if (model.dspark_markov_w1) { + build_dspark_markov_head(*this, model, inp_tokens); + } +} + +// Gemma4-backbone DSpark decoder, dual-mode by batch type (see the DFlash decoder above). +// Gemma4 flavor: sandwich norms, V = the K projection through a scale-less RMS norm (no rope), +// GELU, scaled token embeddings, per-layer output scale and final-logit softcapping +llama_model_dflash::graph_gemma4::graph_gemma4(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + ggml_tensor * inp_pos = build_inp_pos(); + + // optional iSWA: pick the matching attention input + const bool use_iswa = hparams.swa_type != LLAMA_SWA_TYPE_NONE; + + llm_graph_input_attn_kv * inp_attn = nullptr; + llm_graph_input_attn_kv_iswa * inp_attn_iswa = nullptr; + if (use_iswa) { + inp_attn_iswa = build_attn_inp_kv_iswa(); + } else { + inp_attn = build_attn_inp_kv(); + } + + const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + + // KV cache injection + if (ubatch.embd) { + auto inp = std::make_unique(n_embd); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(inp->embd); + + ggml_tensor * inp_g = inp->embd; + cb(inp_g, "inp_g_embeddings", -1); + + res->add_input(std::move(inp)); + + 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 = Kcur; + + 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); + + Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); Kcur = ggml_rope_ext( ctx0, Kcur, inp_pos, layer.rope_freqs, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, @@ -492,7 +682,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm); ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm); - ggml_tensor * Vcur = layer.wv ? build_lora_mm(layer.wv, noise_norm) : Kcur; + ggml_tensor * Vcur = Kcur; Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); @@ -500,10 +690,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - if (layer.wv == nullptr) { - // Gemma4: V is the K projection passed through a scale-less RMS norm, no rope - Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); - } + Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); Qcur = ggml_rope_ext( ctx0, Qcur, inp_pos, layer.rope_freqs, @@ -524,10 +711,8 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il) : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); - if (layer.attn_post_norm) { - cur = build_norm(cur, layer.attn_post_norm, NULL, LLM_NORM_RMS, il); - cb(cur, "attn_post_norm", il); - } + cur = build_norm(cur, layer.attn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_post_norm", il); ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); cb(ffn_inp, "ffn_inp", il); @@ -540,13 +725,11 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra layer.ffn_gate, NULL, layer.ffn_gate_s, layer.ffn_down, NULL, layer.ffn_down_s, NULL, - ffn_act, LLM_FFN_PAR, il); + LLM_FFN_GELU, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); - if (layer.ffn_post_norm) { - cur = build_norm(cur, layer.ffn_post_norm, NULL, LLM_NORM_RMS, il); - cb(cur, "ffn_post_norm", il); - } + cur = build_norm(cur, layer.ffn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_post_norm", il); cur = ggml_add(ctx0, cur, ffn_inp); diff --git a/src/models/models.h b/src/models/models.h index 180b30a46d76..d2465bdbf2b7 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1347,6 +1347,10 @@ struct llama_model_dflash : public llama_model_base { ggml_tensor * build_inp_embd_enc() const; }; + struct graph_gemma4 : public llm_graph_context { + graph_gemma4(const llama_model & model, const llm_graph_params & params); + }; + struct graph_dsv4 : public llama_model_deepseek4::graph { graph_dsv4(const llama_model & model, const llm_graph_params & params); }; From ab118331914e88d07156de945ec880c54d1c99c1 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Wed, 5 Aug 2026 15:33:02 +0800 Subject: [PATCH 06/10] convert: align Gemma4DSparkModel set_vocab with the DSV4 converter Assisted-by: Claude Fable --- conversion/gemma.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/conversion/gemma.py b/conversion/gemma.py index fa454417e22c..5970043f5a66 100644 --- a/conversion/gemma.py +++ b/conversion/gemma.py @@ -858,15 +858,14 @@ class Gemma4DSparkModel(Gemma4Model): def set_vocab(self): if self.target_model_dir is None: - raise ValueError( - "DSpark draft model requires --target-model-dir to be specified. " - "Please provide the path to the target model directory containing the tokenizer." - ) - logger.info(f"DSpark: Using tokenizer from target model: {self.target_model_dir}") + raise ValueError("Gemma4 DSpark requires --target-model-dir with the target tokenizer") + original_dir = self.dir_model - self.dir_model = self.target_model_dir - super().set_vocab() - self.dir_model = original_dir + try: + self.dir_model = self.target_model_dir + super().set_vocab() + finally: + self.dir_model = original_dir mask_token_id = self.hparams.get("mask_token_id") if mask_token_id is not None: From 5e56e841fa38e91d6cfad755d18e0f0f88167740 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Wed, 5 Aug 2026 15:37:44 +0800 Subject: [PATCH 07/10] dspark: shorten the Gemma4 graph comment, generalize the backbone doc Assisted-by: Claude Fable --- docs/speculative.md | 3 +-- src/models/dflash.cpp | 4 +--- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/docs/speculative.md b/docs/speculative.md index 257256524225..06b19e421ffe 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -103,8 +103,7 @@ llama-server -m Qwen3-4B.gguf -md Qwen3-4B-DSpark.gguf \ `--spec-draft-p-min P` also gates the DSpark confidence head: each drafted block is truncated at the first position whose predicted acceptance falls below `P` (default 0 = disabled). -DSpark drafts come with either a Qwen3-style or a Gemma4-style backbone; the backbone is detected -from the checkpoint, and both convert and run the same way. +DSpark drafts support multiple backbones; the backbone is detected from the checkpoint, and all of them convert and run the same way. DSpark drafts exported in the [speculators](https://github.com/vllm-project/speculators) format (for example [`RedHatAI/gemma-4-31B-it-speculator.dspark`](https://huggingface.co/RedHatAI/gemma-4-31B-it-speculator.dspark)) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index a2a732a6364d..0511bb7a7cdc 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -559,9 +559,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra } } -// Gemma4-backbone DSpark decoder, dual-mode by batch type (see the DFlash decoder above). -// Gemma4 flavor: sandwich norms, V = the K projection through a scale-less RMS norm (no rope), -// GELU, scaled token embeddings, per-layer output scale and final-logit softcapping +// Gemma4-backbone DSpark decoder, dual-mode by batch type (see the DFlash decoder above) llama_model_dflash::graph_gemma4::graph_gemma4(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { const int64_t n_embd_head = hparams.n_embd_head_v(); From 7817cabe2766514c960adc1b98d1bf05ea442d52 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Wed, 5 Aug 2026 16:05:06 +0800 Subject: [PATCH 08/10] dspark: load the Gemma4 backbone tensors in a dedicated early-exit branch Assisted-by: Claude Fable --- src/models/dflash.cpp | 45 ++++++++++++++++++++++++++++--------------- 1 file changed, 30 insertions(+), 15 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 0511bb7a7cdc..27ddc74bd2a1 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -173,6 +173,35 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { // optional: reduced-vocab drafts ship their own, full-vocab drafts share the target's via ctx_other output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED); + // Gemma4 backbone: marked by its sandwich norms + if (ml->get_tensor_meta(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", 0).str().c_str())) { + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0); + + // no V projection: V = the K projection + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head_k * n_head }, 0); + layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), { n_embd, n_embd_k_gqa }, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0); + + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0); + layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0); + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0); + layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), { n_embd }, 0); + + layer.out_scale = create_tensor(tn(LLM_TENSOR_LAYER_OUT_SCALE, "weight", i), { 1u }, TENSOR_NOT_REQUIRED); + layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), { n_embd_head_k/2 }, + TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); + } + return; + } + for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; @@ -180,7 +209,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head_k * n_head }, 0); layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), { n_embd, n_embd_k_gqa }, 0); - layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, TENSOR_NOT_REQUIRED); // absent for Gemma4 (V = K projection) + layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, 0); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0); layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0); @@ -190,20 +219,6 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0); layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0); layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0); - - // Gemma4 backbone: sandwich norms, per-layer output scalar and proportional-rope freq factors - layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED); - layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED); - layer.out_scale = create_tensor(tn(LLM_TENSOR_LAYER_OUT_SCALE, "weight", i), { 1u }, TENSOR_NOT_REQUIRED); - layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), { n_embd_head_k/2 }, - TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); - - // the optional tensors must agree on the backbone type: - // Gemma4 has sandwich norms and no V projection (V = K), Qwen has the opposite - const bool gemma = layer.attn_post_norm != nullptr; - if (gemma == (layer.wv != nullptr) || gemma != (layer.ffn_post_norm != nullptr)) { - throw std::runtime_error(format("inconsistent DFlash backbone tensors in layer %d", i)); - } } } From c5171c9e9d3772dce41eb9dd58678dbb02a364c9 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Thu, 6 Aug 2026 15:52:39 +0800 Subject: [PATCH 09/10] convert: slim down the Gemma4 DSpark converter, drop the FFN_PRE_NORM alias The two ambiguous Gemma4 norm names are remapped explicitly instead of going through the FFN_PRE_NORM alias entry, the DeepSpec config checks are dropped, and the sliding-window opt-out is applied only when the draft is all full attention. Assisted-by: Claude Fable --- conversion/gemma.py | 17 ++++++++--------- gguf-py/gguf/constants.py | 1 - 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/conversion/gemma.py b/conversion/gemma.py index 5970043f5a66..b6c1de3b430e 100644 --- a/conversion/gemma.py +++ b/conversion/gemma.py @@ -853,7 +853,7 @@ def set_gguf_parameters(self): @ModelBase.register("Gemma4DSparkModel") class Gemma4DSparkModel(Gemma4Model): - # DSpark draft with a Gemma4 backbone; same DeepSpec flat config schema and shared + # DSpark draft with a Gemma4 backbone model_arch = gguf.MODEL_ARCH.DFLASH def set_vocab(self): @@ -872,13 +872,9 @@ def set_vocab(self): self.gguf_writer.add_mask_token_id(mask_token_id) def set_gguf_parameters(self): - assert not self.hparams.get("enable_moe_block", False), "Gemma4 DSpark with MoE blocks is not supported" - if self.hparams.get("markov_rank", 0) > 0: - assert self.hparams.get("markov_head_type") == "vanilla", "only the vanilla Markov head is supported" - - # the draft uses full attention on every layer - assert all(lt == "full_attention" for lt in self.hparams["layer_types"]) - self.hparams["sliding_window_pattern"] = 1 + # inject Gemma3Model's pattern == 1 sentinel so the sliding_window inherited from the target config is not written + if all(lt == "full_attention" for lt in self.hparams["layer_types"]): + self.hparams["sliding_window_pattern"] = 1 super().set_gguf_parameters() @@ -906,10 +902,13 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca return super().filter_tensors((name, gen)) def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: - # "post_attention_layernorm" is ambiguous in the shared DFlash tensor map -- resolve it explicitly for the Gemma4 backbone + # the shared DFlash tensor map resolves these norm names Qwen-style -- remap them to their Gemma4 meaning if name.endswith(".post_attention_layernorm.weight"): yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_POST_NORM, bid), data_torch) return + if name.endswith(".pre_feedforward_layernorm.weight"): + yield (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_NORM, bid), data_torch) + return yield from super().modify_tensors(data_torch, name, bid) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a81b0b6c142a..f7450fc60958 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -4849,7 +4849,6 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.HC_HEAD_BASE, MODEL_TENSOR.HC_HEAD_SCALE, MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_PRE_NORM, MODEL_TENSOR.FFN_GATE, MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, From 3fdaf0bb368d42bace9f29c57e92ef6f8a61c54e Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Tue, 18 Aug 2026 15:00:02 +0800 Subject: [PATCH 10/10] dspark: fold the Gemma4 draft graph into the shared DFlash decoder Per review feedback, drop the separate graph_gemma4 class and handle the Gemma4 backbone inside graph with runtime branches: V = K projection through a scale-less RMS norm, sandwich norms, GELU, per-layer rope freq factors and the kq_scale/embedding-scale overrides all key off the loaded tensors and hparams. Assisted-by: Claude Fable 5 Co-authored-by: desovo7 <942845546@qq.com> --- src/models/dflash.cpp | 230 +++++------------------------------------- src/models/models.h | 4 - 2 files changed, 26 insertions(+), 208 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 27ddc74bd2a1..3cdedf9c84ef 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -173,8 +173,8 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { // optional: reduced-vocab drafts ship their own, full-vocab drafts share the target's via ctx_other output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED); - // Gemma4 backbone: marked by its sandwich norms - if (ml->get_tensor_meta(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", 0).str().c_str())) { + // Gemma4 backbone: marked by its scaled token embeddings (always written by the Gemma4 converter) + if (hparams.f_embedding_scale != 0.0f) { for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; @@ -231,9 +231,6 @@ std::unique_ptr llama_model_dflash::build_arch_graph(const ll if (hparams.dsv4_hc_mult > 0) { return std::make_unique(*this, params); } - if (layers[0].attn_post_norm) { - return std::make_unique(*this, params); - } return std::make_unique>(*this, params); default: GGML_ABORT("invalid graph type"); @@ -385,197 +382,9 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & // * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens template <> llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { - const int64_t n_embd_head = hparams.n_embd_head_v(); - - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); - - ggml_tensor * inp_pos = build_inp_pos(); - - // optional iSWA: pick the matching attention input - const bool use_iswa = hparams.swa_type != LLAMA_SWA_TYPE_NONE; - - llm_graph_input_attn_kv * inp_attn = nullptr; - llm_graph_input_attn_kv_iswa * inp_attn_iswa = nullptr; - if (use_iswa) { - inp_attn_iswa = build_attn_inp_kv_iswa(); - } else { - inp_attn = build_attn_inp_kv(); - } - - const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); - - // KV cache injection - if (ubatch.embd) { - auto inp = std::make_unique(n_embd); - - inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens); - ggml_set_input(inp->embd); - - ggml_tensor * inp_g = inp->embd; - cb(inp_g, "inp_g_embeddings", -1); - - res->add_input(std::move(inp)); - - 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); - - 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); - - Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); - cb(Kcur, "Kcur_injected", il); - cb(Vcur, "Vcur_injected", il); - - if (use_iswa) { - // route each layer's K/V to its sub-cache: SWA layers -> sliding cache, full -> dense - const bool is_swa = hparams.is_swa(il); - const auto * kv = is_swa ? inp_attn_iswa->mctx->get_swa() : inp_attn_iswa->mctx->get_base(); - ggml_tensor * k_idxs = is_swa ? inp_attn_iswa->get_k_idxs_swa() : inp_attn_iswa->get_k_idxs(); - ggml_tensor * v_idxs = is_swa ? inp_attn_iswa->get_v_idxs_swa() : inp_attn_iswa->get_v_idxs(); - // rotate K/V into the cache's rotated space - ggml_tensor * k_rot = is_swa ? inp_attn_iswa->self_k_rot_swa : inp_attn_iswa->self_k_rot; - ggml_tensor * v_rot = is_swa ? inp_attn_iswa->self_v_rot_swa : inp_attn_iswa->self_v_rot; - if (k_rot) { - Kcur = llama_mul_mat_hadamard(ctx0, Kcur, k_rot); - } - if (v_rot) { - Vcur = llama_mul_mat_hadamard(ctx0, Vcur, v_rot); - } - ggml_build_forward_expand(gf, kv->cpy_k(ctx0, Kcur, k_idxs, il)); - ggml_build_forward_expand(gf, kv->cpy_v(ctx0, Vcur, v_idxs, il)); - } else { - // rotate K/V into the cache's rotated space - if (inp_attn->self_k_rot) { - Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot); - } - if (inp_attn->self_v_rot) { - Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot); - } - ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il)); - ggml_build_forward_expand(gf, inp_attn->mctx->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il)); - } - } - - res->t_embd = inp_g; - - ggml_build_forward_expand(gf, inp_g); - return; - } - - // tok_embd from the target model (shared via ctx_other) - auto * tok_embd = model.tok_embd; - if (tok_embd == nullptr) { - GGML_ASSERT(cparams.ctx_other != nullptr); - const auto * model_other = llama_get_model(cparams.ctx_other); - - GGML_ASSERT(model_other->tok_embd != nullptr && "DFlash decoder requires the target model's token embeddings"); - tok_embd = model_other->tok_embd; - } - - auto inp = std::make_unique(n_embd); - - inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); - ggml_set_input(inp->tokens); - - ggml_tensor * inp_tokens = inp->tokens; - - ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens); - cb(inpL, "inp_noise_embd", -1); - - res->add_input(std::move(inp)); - - for (int il = 0; il < n_layer; ++il) { - const auto & layer = model.layers[il]; - - ggml_tensor * noise_norm = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); - cb(noise_norm, "noise_norm", il); - - ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm); - ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm); - ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm); - - Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); - 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); - - Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); - Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - - Qcur = ggml_rope_ext( - ctx0, Qcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); - cb(Qcur, "Qcur", il); - cb(Kcur, "Kcur", il); - cb(Vcur, "Vcur", il); - - // cache-aware, non-causal attention - ggml_tensor * cur = use_iswa - ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il) - : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + // Gemma4 backbone: marked by its scaled token embeddings (always written by the Gemma4 converter) + const bool is_gemma4 = hparams.f_embedding_scale != 0.0f; - ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); - cb(ffn_inp, "ffn_inp", il); - - cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); - cb(cur, "ffn_norm", il); - - cur = build_ffn(cur, - layer.ffn_up, NULL, NULL, - layer.ffn_gate, NULL, NULL, - layer.ffn_down, NULL, NULL, - NULL, - LLM_FFN_SILU, LLM_FFN_PAR, il); - cb(cur, "ffn_out", il); - - cur = ggml_add(ctx0, cur, ffn_inp); - cb(cur, "l_out", il); - - inpL = cur; - } - - ggml_tensor * cur = build_norm(inpL, model.output_norm, NULL, LLM_NORM_RMS, -1); - cb(cur, "result_norm", -1); - - res->t_embd = cur; - - // lm_head from the target model (shared via ctx_other) - auto * output = model.output; - if (output == nullptr) { - GGML_ASSERT(cparams.ctx_other != nullptr); - const auto * model_other = llama_get_model(cparams.ctx_other); - GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection"); - output = model_other->output; - } - - cur = build_lora_mm(output, cur); - cb(cur, "result_output", -1); - res->t_logits = cur; - - ggml_build_forward_expand(gf, cur); - - // DSpark: bias the draft logits with the Markov head - if (model.dspark_markov_w1) { - build_dspark_markov_head(*this, model, inp_tokens); - } -} - -// Gemma4-backbone DSpark decoder, dual-mode by batch type (see the DFlash decoder above) -llama_model_dflash::graph_gemma4::graph_gemma4(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { const int64_t n_embd_head = hparams.n_embd_head_v(); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); @@ -595,6 +404,9 @@ llama_model_dflash::graph_gemma4::graph_gemma4(const llama_model & model, const const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + // the Gemma4 backbone uses GELU instead of SiLU + const llm_ffn_op_type ffn_act = is_gemma4 ? LLM_FFN_GELU : LLM_FFN_SILU; + // KV cache injection if (ubatch.embd) { auto inp = std::make_unique(n_embd); @@ -611,13 +423,16 @@ llama_model_dflash::graph_gemma4::graph_gemma4(const llama_model & model, const const auto & layer = model.layers[il]; ggml_tensor * Kcur = build_lora_mm(layer.wk, inp_g); - ggml_tensor * Vcur = Kcur; + // Gemma4 has no V projection: V = the K projection through a scale-less RMS norm + ggml_tensor * Vcur = is_gemma4 ? Kcur : build_lora_mm(layer.wv, inp_g); 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); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); + if (is_gemma4) { + Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); + } Kcur = ggml_rope_ext( ctx0, Kcur, inp_pos, layer.rope_freqs, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, @@ -695,7 +510,8 @@ llama_model_dflash::graph_gemma4::graph_gemma4(const llama_model & model, const ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm); ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm); - ggml_tensor * Vcur = Kcur; + // Gemma4 has no V projection: V = the K projection through a scale-less RMS norm + ggml_tensor * Vcur = is_gemma4 ? Kcur : build_lora_mm(layer.wv, noise_norm); Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); @@ -703,7 +519,9 @@ llama_model_dflash::graph_gemma4::graph_gemma4(const llama_model & model, const Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); + if (is_gemma4) { + Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); + } Qcur = ggml_rope_ext( ctx0, Qcur, inp_pos, layer.rope_freqs, @@ -724,8 +542,10 @@ llama_model_dflash::graph_gemma4::graph_gemma4(const llama_model & model, const ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il) : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); - cur = build_norm(cur, layer.attn_post_norm, NULL, LLM_NORM_RMS, il); - cb(cur, "attn_post_norm", il); + if (layer.attn_post_norm) { + cur = build_norm(cur, layer.attn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_post_norm", il); + } ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); cb(ffn_inp, "ffn_inp", il); @@ -738,11 +558,13 @@ llama_model_dflash::graph_gemma4::graph_gemma4(const llama_model & model, const layer.ffn_gate, NULL, layer.ffn_gate_s, layer.ffn_down, NULL, layer.ffn_down_s, NULL, - LLM_FFN_GELU, LLM_FFN_PAR, il); + ffn_act, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); - cur = build_norm(cur, layer.ffn_post_norm, NULL, LLM_NORM_RMS, il); - cb(cur, "ffn_post_norm", il); + if (layer.ffn_post_norm) { + cur = build_norm(cur, layer.ffn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_post_norm", il); + } cur = ggml_add(ctx0, cur, ffn_inp); diff --git a/src/models/models.h b/src/models/models.h index d2465bdbf2b7..180b30a46d76 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1347,10 +1347,6 @@ struct llama_model_dflash : public llama_model_base { ggml_tensor * build_inp_embd_enc() const; }; - struct graph_gemma4 : public llm_graph_context { - graph_gemma4(const llama_model & model, const llm_graph_params & params); - }; - struct graph_dsv4 : public llama_model_deepseek4::graph { graph_dsv4(const llama_model & model, const llm_graph_params & params); };