download: find MTP heads published in a dedicated MTP/ folder

common/download.cpp

@@ -17,6 +17,7 @@
 #include <unordered_set>
 #include <string>
 #include <thread>
+#include <tuple>
 #include <vector>
 
 #include "http.h"
@@ -563,23 +564,40 @@ static hf_cache::hf_files get_split_files(const hf_cache::hf_files & files,
     return result;
 }
 
+static std::string to_upper(const std::string & s) {
+    std::string r = s;
+    for (char & c : r) {
+        c = (char) std::toupper((unsigned char) c);
+    }
+    return r;
+}
+
 // pick the best sibling GGUF whose filename contains `keyword` (e.g. "mmproj" / "mtp"),
 // preferring deeper shared directory prefix with the model, then exact `tag` match,
 // then closest quantization to the tag when given, or to the model otherwise
+//
+// `sidecar_dir` additionally admits candidates in a dedicated top-level folder of that name, which
+// shares no directory prefix with the model and is otherwise unreachable: unsloth publishes MTP
+// heads as MTP/mtp-<model>-<quant>.gguf, and without this the folder is silently never matched
+// `prefer_bits` wins over proximity when no candidate matches `tag` exactly, so a repo offering
+// several heads has one designated default rather than whichever is arithmetically nearest
+//
+// both only ever apply to candidates admitted by `sidecar_dir`, so a repository that resolved
+// before this existed resolves to the same file afterwards
 static hf_cache::hf_file find_best_sibling(const hf_cache::hf_files & files,
                                            const std::string        & model,
                                            const std::string        & keyword,
-                                           const std::string        & tag = "") {
+                                           const std::string        & tag = "",
+                                           const std::string        & sidecar_dir = "",
+                                           int                        prefer_bits = 0) {
+    using rank_t = std::tuple<bool, size_t, bool, bool, int, bool>;
+
     hf_cache::hf_file best;
-    size_t best_depth = 0;
-    int best_diff = 0;
-    bool best_exact = false;
+    rank_t best_rank;
     bool found = false;
 
-    std::string tag_upper = tag;
-    for (char & c : tag_upper) {
-        c = (char) std::toupper((unsigned char) c);
-    }
+    const std::string tag_upper     = to_upper(tag);
+    const std::string sidecar_upper = to_upper(sidecar_dir);
 
     int model_bits = 0;
     if (!tag_upper.empty()) {
@@ -602,27 +620,35 @@ static hf_cache::hf_file find_best_sibling(const hf_cache::hf_files & files,
 
         auto [_, dir] = std::mismatch(model_parts.begin(), model_dir,
                                       sib_parts.begin(), sib_dir);
-        if (dir != sib_dir) {
+
+        size_t depth  = 0;
+        bool   prefix = dir == sib_dir;
+        if (prefix) {
+            depth = dir - sib_parts.begin();
+        } else if (sidecar_upper.empty() || sib_parts.size() != 2 ||
+                   to_upper(sib_parts[0]) != sidecar_upper) {
             continue;
         }
 
-        size_t depth = dir - sib_parts.begin();
         auto bits = extract_quant_bits(f.path);
         auto diff = std::abs(bits - model_bits);
 
-        std::string path_upper = f.path;
-        for (char & c : path_upper) {
-            c = (char) std::toupper((unsigned char) c);
-        }
+        const std::string path_upper = to_upper(f.path);
+
         bool exact = !tag_upper.empty() && path_upper.find("-" + tag_upper + ".") != std::string::npos;
 
-        if (!found || depth > best_depth ||
-            (depth == best_depth && exact && !best_exact) ||
-            (depth == best_depth && exact == best_exact && diff < best_diff)) {
+        // a sibling sharing the model's directory prefix always outranks the dedicated folder, and
+        // the two rules below never fire for one, so every layout that resolved before is unchanged
+        bool preferred = !prefix && prefer_bits != 0 && bits == prefer_bits;
+        // a head carrying its own token_embd/output loads against any build, where one that borrows
+        // them from the target needs a build that supports the borrow
+        bool self_contained = !prefix && path_upper.find("SHARED") == std::string::npos;
+
+        rank_t rank{prefix, depth, exact, preferred, -diff, self_contained};
+
+        if (!found || rank > best_rank) {
             best = f;
-            best_depth = depth;
-            best_diff = diff;
-            best_exact = exact;
+            best_rank = rank;
             found = true;
         }
     }
@@ -634,10 +660,12 @@ static hf_cache::hf_file find_best_mmproj(const hf_cache::hf_files & files,
     return find_best_sibling(files, model, "mmproj");
 }
 
+// 8 bits is Q8_0: measured the fastest MTP head as well as the most accurate, because a draft step
+// is dominated by the LM head and that head is cheaper to execute at 8 bits than at bf16
 static hf_cache::hf_file find_best_mtp(const hf_cache::hf_files & files,
                                        const std::string        & model,
                                        const std::string        & tag = "") {
-    return find_best_sibling(files, model, "mtp-", tag);
+    return find_best_sibling(files, model, "mtp-", tag, "MTP", 8);
 }
 
 static hf_cache::hf_file find_best_eagle3(const hf_cache::hf_files & files,

tests/test-model-resolution.cpp

@@ -180,6 +180,41 @@ static const std::vector<std::string> spark = {
     "dspark-model-MXFP4.gguf",
 };
 
+// heads in a dedicated folder beside a root model, in the style of unsloth/Qwen3.8-27B-GGUF
+static const std::vector<std::string> mtp_dir_flat = {
+    "model-Q4_K_M.gguf",
+    "model-Q8_0.gguf",
+    "MTP/mtp-model-Q4_0.gguf",
+};
+
+// the same folder beside quant subdirectories, in the style of unsloth/Qwen3.8-Flash-Next-GGUF,
+// carrying every head we publish: three quants, each self-contained and borrowing
+static const std::vector<std::string> mtp_dir_subdir = {
+    "mmproj-BF16.gguf",
+    "UD-IQ1_S/model-UD-IQ1_S-00001-of-00002.gguf",
+    "UD-IQ1_S/model-UD-IQ1_S-00002-of-00002.gguf",
+    "Q8_0/model-Q8_0.gguf",
+    "MTP/mtp-model-BF16.gguf",
+    "MTP/mtp-model-Q4_K_M.gguf",
+    "MTP/mtp-model-Q8_0.gguf",
+    "MTP/mtp-model-shared-BF16.gguf",
+    "MTP/mtp-model-shared-Q4_K_M.gguf",
+    "MTP/mtp-model-shared-Q8_0.gguf",
+};
+
+// a dedicated folder must not take precedence over a head sharing the model's own directory
+static const std::vector<std::string> mtp_dir_and_sibling = {
+    "model-Q8_0.gguf",
+    "mtp-model-BF16.gguf",
+    "MTP/mtp-model-Q8_0.gguf",
+};
+
+// a top-level name that merely starts with MTP is not the dedicated folder
+static const std::vector<std::string> mtp_dir_near_miss = {
+    "model-Q8_0.gguf",
+    "MTPX/mtp-model-Q8_0.gguf",
+};
+
 // dspark outranks dflash in the type auto-selection
 static const std::vector<std::string> dspark_dflash = {
     "model-Q8_0.gguf",
@@ -283,6 +318,37 @@ static const plan_case plan_cases[] = {
     {"spark tag sidecar", spark, "test/repo:BF16", "", true, false,
      "", {},
      "", "", "", "", "dspark-model-BF16.gguf"},
+
+    // a head in the dedicated folder is reachable from a model at the repo root,
+    // where the directory prefix rule alone finds nothing
+    {"mtp dir flat", mtp_dir_flat, "test/repo:Q8_0", "", true, false,
+     "model-Q8_0.gguf", {"model-Q8_0.gguf"},
+     "", "MTP/mtp-model-Q4_0.gguf", "", "", ""},
+
+    // and from a model in a quant subdirectory, which shares no prefix with it either.
+    // no head matches UD-IQ1_S, so the preferred quant wins over the nearest one, and the
+    // self-contained head wins over the one that borrows from the target
+    {"mtp dir subdir", mtp_dir_subdir, "test/repo:UD-IQ1_S", "", true, false,
+     "UD-IQ1_S/model-UD-IQ1_S-00001-of-00002.gguf",
+     {"UD-IQ1_S/model-UD-IQ1_S-00001-of-00002.gguf",
+      "UD-IQ1_S/model-UD-IQ1_S-00002-of-00002.gguf"},
+     "mmproj-BF16.gguf", "MTP/mtp-model-Q8_0.gguf", "", "", ""},
+
+    // an exact tag still beats the preferred quant
+    {"mtp dir exact tag", mtp_dir_subdir, "test/repo:Q8_0", "", true, false,
+     "Q8_0/model-Q8_0.gguf", {"Q8_0/model-Q8_0.gguf"},
+     "mmproj-BF16.gguf", "MTP/mtp-model-Q8_0.gguf", "", "", ""},
+
+    // a head sharing the model's own directory outranks the dedicated folder, even though the
+    // folder holds the exact tag and the sibling does not
+    {"mtp dir loses to sibling", mtp_dir_and_sibling, "test/repo:Q8_0", "", true, false,
+     "model-Q8_0.gguf", {"model-Q8_0.gguf"},
+     "", "mtp-model-BF16.gguf", "", "", ""},
+
+    // only the exact folder name is special, MTPX/ is not
+    {"mtp dir near miss", mtp_dir_near_miss, "test/repo:Q8_0", "", true, false,
+     "model-Q8_0.gguf", {"model-Q8_0.gguf"},
+     "", "", "", "", ""},
 };
 
 static void check_plan(const plan_case & c) {