From 29a54f4b04cc0281850574247db075e20d68e7f1 Mon Sep 17 00:00:00 2001 From: Joel Farthing Date: Mon, 29 Jun 2026 06:26:29 -0500 Subject: [PATCH] DFlash: support MiMo-V2.5-Pro draft conversion and runtime (#2048) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Support MiMo DFlash draft conversion * Fix MiMo2 DFlash capture row pruning * Fix MiMo DFlash draft RoPE and value scale * Honor partial_rotary_factor in DFlash draft RoPE dim count The draft set rope.dimension_count to the full head_dim (128), ignoring the MiMo DFlash draft's partial_rotary_factor=0.5. The correct count is head_dim*partial_rotary_factor=64; the remaining dims are NoPE. With the full head_dim the upper half of each head receives position rotation it was never trained for, which roughly halves draft acceptance on code (~26% -> ~60% once corrected). RoPE base (5e6) and value scale (0.612) were already correct. * Filter weight-map shard discovery to files that exist get_model_part_names_from_weight_map() returned shard names straight from the index weight_map without checking they exist. A model dir with a stale model.safetensors.index.json but no safetensors shards would then set is_safetensors=True and skip the pytorch_model*.bin fallback, failing later when opening the missing files. Filter to shards present on disk so a stale index falls through to the other weight formats. * DFlash: store backbone_rotary_base in dedicated GGUF key backbone_rotary_base (the target model's RoPE theta used when encoding context K/V) was written to rope.freq_base, clobbering the draft model's own rope_theta. For MiMo this swapped 10000 → 5000000 in the draft attention path. Fix: write backbone_rotary_base to a dedicated dflash.backbone_rotary_base GGUF key and read it into hparams.dflash_backbone_rotary_base. In build_dflash_kv_cache, use target_freq_base (the new hparam when set, falling back to freq_base) for the context-K RoPE call. The draft model's own rope.freq_base is now set correctly from rope_theta. Existing MiMo DFlash GGUFs must be reconverted. --------- Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com> --- convert_hf_to_gguf.py | 47 +++++++++++++++++++++++++++++++++- gguf-py/gguf/constants.py | 5 ++++ gguf-py/gguf/gguf_writer.py | 3 +++ gguf-py/gguf/tensor_mapping.py | 4 +++ src/graphs/build_dflash.cpp | 15 ++++++++++- src/graphs/build_mimo2.cpp | 12 +++++++-- src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/llama-hparams.cpp | 8 +++--- src/llama-hparams.h | 2 ++ src/llama-load-tensors.cpp | 1 + src/llama-model.cpp | 1 + 12 files changed, 93 insertions(+), 7 deletions(-) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 31ba6e1e..a9c40c8c 100644 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -85,9 +85,13 @@ class Model: self.use_temp_file = use_temp_file self.lazy = not eager self.part_names = Model.get_model_part_names(self.dir_model, "model", ".safetensors") + if len(self.part_names) == 0: + self.part_names = Model.get_model_part_names_from_weight_map(self.dir_model, "model.safetensors.index.json") self.is_safetensors = len(self.part_names) > 0 if not self.is_safetensors: self.part_names = Model.get_model_part_names(self.dir_model, "pytorch_model", ".bin") + if len(self.part_names) == 0: + self.part_names = Model.get_model_part_names_from_weight_map(self.dir_model, "pytorch_model.bin.index.json") self.hparams = Model.load_hparams(self.dir_model) self.block_count = self.find_hparam(["n_layers", "num_hidden_layers", "n_layer", "num_layers"]) self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) @@ -457,6 +461,24 @@ class Model: return part_names + @staticmethod + def get_model_part_names_from_weight_map(dir_model: Path, index_name: str) -> list[str]: + index_path = dir_model / index_name + if not index_path.exists(): + return [] + + with open(index_path, "r", encoding="utf-8") as f: + index: dict[str, Any] = json.load(f) + weight_map = index.get("weight_map") + if weight_map is None or not isinstance(weight_map, dict): + raise ValueError(f"Can't load 'weight_map' from {index_name!r}") + + part_names = sorted({str(part_name) for part_name in weight_map.values()}) + # Only surface shards that exist on disk; a stale index.json would otherwise set + # is_safetensors=True and suppress the pytorch_model*.bin fallback. + part_names = [name for name in part_names if (dir_model / name).is_file()] + return part_names + @staticmethod def load_hparams(dir_model: Path): with open(dir_model / "config.json", "r", encoding="utf-8") as f: @@ -2399,7 +2421,13 @@ class DFlashDraftModel(Qwen3Model): super().set_gguf_parameters() self.gguf_writer.add_causal_attention(False) - self.gguf_writer.add_rope_dimension_count(self.hparams.get("head_dim", 128)) + # MiMo DFlash draft uses partial rotary (partial_rotary_factor=0.5): RoPE is applied to + # only head_dim*partial_rotary_factor dims, the rest are NoPE. Honoring it is required; + # otherwise the upper half of every head gets spurious position rotation it was never + # trained for, which roughly halves draft acceptance. + head_dim = self.hparams.get("head_dim", 128) + partial_rotary_factor = self.hparams.get("partial_rotary_factor", 1.0) + self.gguf_writer.add_rope_dimension_count(int(partial_rotary_factor * head_dim)) rope_scaling = self.hparams.get("rope_scaling") if isinstance(rope_scaling, dict): @@ -2428,6 +2456,15 @@ class DFlashDraftModel(Qwen3Model): dflash_cfg = self.hparams.get("dflash_config") dflash_cfg = dflash_cfg if isinstance(dflash_cfg, dict) else {} + if (backbone_rotary_base := dflash_cfg.get("backbone_rotary_base")) is not None: + self.gguf_writer.add_float32(f"{arch}.dflash.backbone_rotary_base", float(backbone_rotary_base)) + logger.info("DFlashDraftModel: backbone_rotary_base=%s", backbone_rotary_base) + + attention_value_scale = dflash_cfg.get("attention_value_scale", self.hparams.get("attention_value_scale")) + if attention_value_scale is not None: + self.gguf_writer.add_attention_value_scale(float(attention_value_scale)) + logger.info("DFlashDraftModel: attention_value_scale=%s", attention_value_scale) + def dflash_required_value(name: str) -> Any: if name in dflash_cfg: return dflash_cfg[name] @@ -2484,6 +2521,10 @@ class DFlashDraftModel(Qwen3Model): # all-SWA only when it is absent. Absent/false use_sliding_window => dense draft (unchanged). use_sliding_window = self.hparams.get("use_sliding_window") sliding_window = self.hparams.get("sliding_window") + if use_sliding_window is None and "use_swa" in dflash_cfg: + use_sliding_window = bool(dflash_cfg["use_swa"]) + if sliding_window is None and "swa_window_size" in dflash_cfg: + sliding_window = int(dflash_cfg["swa_window_size"]) if use_sliding_window and sliding_window: n_swa_layers = int(self.hparams.get("num_hidden_layers", self.block_count)) layer_types = self.hparams.get("layer_types") @@ -2542,6 +2583,10 @@ class DFlashDraftModel(Qwen3Model): return [(f"{gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DFLASH_FC]}.weight", data_torch)] if top_level_name == "hidden_norm.weight": return [(f"{gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DFLASH_HIDDEN_NORM]}.weight", data_torch)] + if top_level_name.endswith(".self_attn.attention_sink_bias"): + if bid is None: + raise ValueError(f"DFlashDraftModel: can not infer block id for tensor {name!r}") + return [(f"{gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.ATTN_SINKS].format(bid=bid)}.weight", data_torch)] if name == "norm.weight": name = "model.norm.weight" elif name.startswith("layers."): diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 76c81d9e..e3b05643 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -119,6 +119,7 @@ class Keys: SHARED_KV_LAYERS = "{arch}.attention.shared_kv_layers" KEY_LENGTH_SWA = "{arch}.attention.key_length_swa" VALUE_LENGTH_SWA = "{arch}.attention.value_length_swa" + VALUE_SCALE = "{arch}.attention.value_scale" OUTPUT_SCALE = "{arch}.attention.output_scale" TEMPERATURE_LENGTH = "{arch}.attention.temperature_length" @@ -319,6 +320,7 @@ class MODEL_TENSOR(IntEnum): FFN_EXP_PROBS_B = auto() ATTN_Q_NORM = auto() ATTN_K_NORM = auto() + ATTN_SINKS = auto() LAYER_OUT_NORM = auto() LAYER_OUT_SCALE = auto() PER_LAYER_TOKEN_EMBD = auto() @@ -470,6 +472,7 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.ATTN_ROT_EMBD: "blk.{bid}.attn_rot_embd", MODEL_TENSOR.ATTN_Q_NORM: "blk.{bid}.attn_q_norm", MODEL_TENSOR.ATTN_K_NORM: "blk.{bid}.attn_k_norm", + MODEL_TENSOR.ATTN_SINKS: "blk.{bid}.attn_sinks", MODEL_TENSOR.ATTN_OUT_NORM: "blk.{bid}.attn_output_norm", MODEL_TENSOR.ATTN_POST_NORM: "blk.{bid}.post_attention_norm", MODEL_TENSOR.ATTN_GATE: "blk.{bid}.attn_gate", @@ -1303,6 +1306,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_K_NORM, MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_SINKS, MODEL_TENSOR.ATTN_OUT, MODEL_TENSOR.ATTN_POST_NORM, MODEL_TENSOR.FFN_GATE, @@ -1321,6 +1325,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_K_NORM, MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_SINKS, MODEL_TENSOR.ATTN_OUT, MODEL_TENSOR.ATTN_POST_NORM, MODEL_TENSOR.FFN_GATE, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 68440f0d..09fc6aba 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -851,6 +851,9 @@ class GGUFWriter: def add_attention_scale(self, value: float) -> None: self.add_float32(Keys.Attention.SCALE.format(arch=self.arch), value) + def add_attention_value_scale(self, value: float) -> None: + self.add_float32(Keys.Attention.VALUE_SCALE.format(arch=self.arch), value) + def add_attn_output_scale(self, value: float) -> None: self.add_float32(Keys.Attention.OUTPUT_SCALE.format(arch=self.arch), value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 63932395..2a5f3841 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -423,6 +423,10 @@ class TensorNameMap: "model.layers.{bid}.attention.key_layernorm", # bailingmoe2 ), + MODEL_TENSOR.ATTN_SINKS: ( + "model.layers.{bid}.self_attn.attention_sink_bias", # MiMo DFlash + ), + MODEL_TENSOR.ROPE_FREQS: ( "language_model.encoder.layers.{bid}.self_attention.rotary_emb.inv_freq", # persimmon ), diff --git a/src/graphs/build_dflash.cpp b/src/graphs/build_dflash.cpp index 725d08e5..4f6c7cd0 100644 --- a/src/graphs/build_dflash.cpp +++ b/src/graphs/build_dflash.cpp @@ -44,8 +44,10 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() { ggml_tensor * Kcur_ctx = ggml_reshape_3d(ctx0, Kcur_ctx_proj, n_embd_head_k, n_head_kv, update_rows); Kcur_ctx = llm_build_norm(ctx0, Kcur_ctx, hparams, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, cb, il); cb(Kcur_ctx, "dflash_kv_k_norm", il); + const float target_freq_base = hparams.dflash_backbone_rotary_base > 0.0f + ? hparams.dflash_backbone_rotary_base : freq_base; Kcur_ctx = ggml_rope_ext(ctx0, Kcur_ctx, lctx.dflash.kv.cache_input_pos_ctx, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + n_rot, rope_type, n_ctx_orig, target_freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); cb(Kcur_ctx, "dflash_kv_k_rope", il); Kcur_ctx = ggml_cont(ctx0, ggml_permute(ctx0, Kcur_ctx, 0, 2, 1, 3)); @@ -53,6 +55,10 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() { ggml_tensor * Vcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, fused_target); cb(Vcur_ctx, "dflash_kv_v_proj", il); + if (std::abs(hparams.f_attn_v_scale - 1.0f) > 1e-4f) { + Vcur_ctx = ggml_scale(ctx0, Vcur_ctx, hparams.f_attn_v_scale); + cb(Vcur_ctx, "dflash_kv_v_scaled", il); + } Vcur_ctx = ggml_reshape_3d(ctx0, Vcur_ctx, n_embd_head_v, n_head_kv, update_rows); Vcur_ctx = ggml_cont(ctx0, ggml_permute(ctx0, Vcur_ctx, 0, 2, 1, 3)); cb(Vcur_ctx, "dflash_kv_v_physical", il); @@ -252,6 +258,10 @@ ggml_cgraph * llm_build_context::build_dflash() { cb(Kcur_noise, "Kcur_roped", il); Vcur_noise = ggml_reshape_3d(ctx0, Vcur_noise, n_embd_head_v, n_head_kv, n_tokens); + if (std::abs(hparams.f_attn_v_scale - 1.0f) > 1e-4f) { + Vcur_noise = ggml_scale(ctx0, Vcur_noise, hparams.f_attn_v_scale); + cb(Vcur_noise, "Vcur_noise_scaled", il); + } cb(Vcur_noise, "Vcur_noise", il); GGML_ASSERT(il < (int32_t) lctx.dflash.kv.k_ctx_cache.size()); @@ -313,6 +323,9 @@ ggml_cgraph * llm_build_context::build_dflash() { cur = ggml_flash_attn_ext(ctx0, q, k, v, dflash_kq_mask_l, kq_scale, hparams.f_max_alibi_bias, hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); + if (model.layers[il].attn_sinks) { + ggml_flash_attn_ext_add_sinks(cur, model.layers[il].attn_sinks); + } cb(cur, "flash_attn", il); ggml_build_forward_expand(gf, cur); // Somethiong goes wrong with thisi optimization. diff --git a/src/graphs/build_mimo2.cpp b/src/graphs/build_mimo2.cpp index cb5a1ea1..78f15164 100644 --- a/src/graphs/build_mimo2.cpp +++ b/src/graphs/build_mimo2.cpp @@ -17,6 +17,7 @@ ggml_cgraph * llm_build_context::build_mimo2() { // inp_pos - contains the positions struct ggml_tensor * inp_pos = build_inp_pos(); struct ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr; + const bool needs_dflash_full_final_rows = lctx.dflash.capture != nullptr; // KQ_mask (mask for 1 head, it will be broadcasted to all heads) struct ggml_tensor * KQ_mask = build_inp_KQ_mask(); @@ -26,8 +27,11 @@ ggml_cgraph * llm_build_context::build_mimo2() { const bool is_sliding = model.hparams.swa_layers[il]; auto KQ_mask_l = is_sliding ? KQ_mask_swa : KQ_mask; + const bool is_final_layer = il == n_layer - 1; + struct ggml_tensor * attn_out_ids = is_final_layer && !needs_dflash_full_final_rows ? inp_out_ids : nullptr; + cur = build_std_attention(gf, model.layers[il].attn_norm, inpL, - inp_pos, il == n_layer - 1 ? inp_out_ids : nullptr, nullptr, + inp_pos, attn_out_ids, nullptr, KQ_mask_l, model.layers[il].attn_sinks, nullptr, 1.0f/sqrtf(float(n_embd_head_k)), 0.0f, is_sliding ? hparams.n_swa : 0, il, true, false, true); @@ -60,6 +64,11 @@ ggml_cgraph * llm_build_context::build_mimo2() { cur = lctx.cvec.apply_to(ctx0, cur, il); cb(cur, "l_out", il); + if (is_final_layer && needs_dflash_full_final_rows && inp_out_ids != nullptr) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + cb(cur, "l_out_selected", il); + } + // input for next layer inpL = cur; } @@ -73,4 +82,3 @@ ggml_cgraph * llm_build_context::build_mimo2() { return gf; } - diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 513acd55..46f6882b 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -165,6 +165,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_DFLASH_MASK_TOKEN_ID, "%s.dflash.mask_token_id" }, { LLM_KV_DFLASH_TARGET_LAYER_IDS, "%s.dflash.target_layer_ids" }, { LLM_KV_DFLASH_N_TARGET_FEATURES, "%s.dflash.n_target_features" }, + { LLM_KV_DFLASH_BACKBONE_ROTARY_BASE, "%s.dflash.backbone_rotary_base" }, { LLM_KV_ATTENTION_HEAD_COUNT, "%s.attention.head_count" }, { LLM_KV_ATTENTION_HEAD_COUNT_KV, "%s.attention.head_count_kv" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index db37a2ba..1bb21d12 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -148,6 +148,7 @@ enum llm_kv { LLM_KV_DFLASH_MASK_TOKEN_ID, LLM_KV_DFLASH_TARGET_LAYER_IDS, LLM_KV_DFLASH_N_TARGET_FEATURES, + LLM_KV_DFLASH_BACKBONE_ROTARY_BASE, LLM_KV_ATTENTION_HEAD_COUNT, LLM_KV_ATTENTION_HEAD_COUNT_KV, diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index bfb93283..f3fcadf8 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -899,10 +899,12 @@ void llm_load_hparams( case LLM_ARCH_DFLASH_DRAFT: { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); - ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE, hparams.dflash_block_size, false); - ml.get_key(LLM_KV_DFLASH_MASK_TOKEN_ID, hparams.dflash_mask_token_id, false); - ml.get_key(LLM_KV_DFLASH_N_TARGET_FEATURES, hparams.dflash_n_target_features, false); + ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE, hparams.dflash_block_size, false); + ml.get_key(LLM_KV_DFLASH_MASK_TOKEN_ID, hparams.dflash_mask_token_id, false); + ml.get_key(LLM_KV_DFLASH_N_TARGET_FEATURES, hparams.dflash_n_target_features, false); + ml.get_key(LLM_KV_DFLASH_BACKBONE_ROTARY_BASE, hparams.dflash_backbone_rotary_base, false); load_dflash_target_layer_ids(ml, LLM_KV(model.arch)(LLM_KV_DFLASH_TARGET_LAYER_IDS), hparams, false); + ml.get_key(LLM_KV_ATTENTION_VALUE_SCALE, hparams.f_attn_v_scale, false); // DFlash drafts may be trained with sliding-window attention (for long-context). // Read the window + per-layer pattern so the SWA mask path activates; absent keys // leave n_swa=0 / swa_layers all-zero (dense behavior, unchanged). diff --git a/src/llama-hparams.h b/src/llama-hparams.h index a26bd68b..941c75f9 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -145,6 +145,7 @@ struct llama_hparams { uint32_t dflash_n_target_features = 0; uint32_t dflash_n_target_layers = 0; uint32_t dflash_target_layer_ids[8] = {}; + float dflash_backbone_rotary_base = 0.0f; // needed by encoder-decoder models (e.g. T5, FLAN-T5) // ref: https://github.com/ggerganov/llama.cpp/pull/8141 @@ -217,6 +218,7 @@ struct llama_hparams { if (!is_float_close(this->f_residual_scale, other.f_residual_scale, EPSILON)) return true; if (!is_float_close(this->f_embedding_scale, other.f_embedding_scale, EPSILON)) return true; if (!is_float_close(this->f_attention_scale, other.f_attention_scale, EPSILON)) return true; + if (!is_float_close(this->f_attn_v_scale, other.f_attn_v_scale, EPSILON)) return true; return false; } diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index a1300cce..0e1464ee 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -2314,6 +2314,7 @@ bool create_tensors_helper::create_dflash_tensors(const LLM_TN & tn) { layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_sinks = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, llama_model_loader::TENSOR_NOT_REQUIRED); layer.ffn_gate = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); layer.ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index e8aa5be7..09342718 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -857,6 +857,7 @@ static const std::map> LLM_TENSOR_NA { LLM_TENSOR_ATTN_K, "blk.%d.attn_k" }, { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, { LLM_TENSOR_ATTN_V, "blk.%d.attn_v" }, + { LLM_TENSOR_ATTN_SINKS, "blk.%d.attn_sinks" }, { LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" }, { LLM_TENSOR_ATTN_POST_NORM, "blk.%d.post_attention_norm" }, { LLM_TENSOR_FFN_GATE, "blk.%d.ffn_gate" },