DFlash: support MiMo-V2.5-Pro draft conversion and runtime (#2048)

* 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

@@ -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."):

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,

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)
 

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

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.

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

src/llama-arch.cpp

@@ -165,6 +165,7 @@ static const std::map<llm_kv, const char *> 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"          },

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,

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

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

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);

src/llama-model.cpp

@@ -857,6 +857,7 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> 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" },