From 05932db5ec743673bc754873d6a596b421a8af56 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Sun, 26 Jul 2026 23:29:37 -0700
Subject: [PATCH] Studio: share one quant vocabulary with the companion search
 root

The review point was correct. _local_gguf_companion_search_root carried
its own copy of the quant pattern, which never gained the bpw modifier,
so a directory such as IQ4_XS-3.53bpw was not recognised as a quant dir.
The search root stayed inside it and the repository-root MTP/ copy was
out of scope for discovery, the native fallback and reload dedup alike.

It now builds on _GGUF_KNOWN_QUANT_RE plus the optional bpw suffix, so
there is a single vocabulary rather than a duplicate that can fall
behind again. Promotion is unchanged for every previously matching name
and still rejects DeepSeek-V3-UD-Q2_K_XL, Q4_0-extra and Q4_0bpw.
---
 .../tests/test_mtp_drafter_companion.py       | 28 +++++++++++++++++++
 studio/backend/utils/models/model_config.py   | 14 ++--------
 2 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py
index bf1ef1a061..95179b86bd 100644
--- a/studio/backend/tests/test_mtp_drafter_companion.py
+++ b/studio/backend/tests/test_mtp_drafter_companion.py
@@ -799,3 +799,31 @@ def test_detect_mtp_file_ranks_split_drafter_by_total_size(tmp_path):
     smaller.write_bytes(b"x" * 100)
 
     assert detect_mtp_file(str(weight)) == str(smaller.resolve())
+
+
+def test_companion_search_root_promotes_bpw_quant_directory(tmp_path):
+    """A bpw-qualified quant directory must resolve to the repository root, or
+    the repo-root MTP/ copy is never in scope for it."""
+    quant_dir = tmp_path / "IQ4_XS-3.53bpw"
+    quant_dir.mkdir()
+    weight = quant_dir / "model.gguf"
+    weight.write_bytes(b"x")
+    sub = tmp_path / "MTP"
+    sub.mkdir()
+    drafter = sub / "mtp-model.gguf"
+    drafter.write_bytes(b"x")
+
+    # Directory selection and the file inside it agree on the root.
+    assert _local_gguf_companion_search_root(str(quant_dir), str(weight)) == str(tmp_path)
+    assert _local_gguf_companion_search_root(str(weight), str(weight)) == str(tmp_path)
+    assert detect_mtp_file(str(weight), str(tmp_path)) == str(drafter.resolve())
+
+
+def test_companion_search_root_keeps_non_quant_directories(tmp_path):
+    """Sharing the quant vocabulary must not widen what gets promoted."""
+    for name in ("DeepSeek-V3-UD-Q2_K_XL", "outputs", "Q4_0-extra", "Q4_0bpw"):
+        directory = tmp_path / name
+        directory.mkdir()
+        weight = directory / "model.gguf"
+        weight.write_bytes(b"x")
+        assert _local_gguf_companion_search_root(str(directory), str(weight)) == str(directory)
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index 523b2a7b30..e6c1d0b4d7 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -1865,17 +1865,9 @@ def _local_gguf_companion_search_root(selected_path: str, gguf_file: str) -> str
 
     selected = Path(selected_path)
     gguf_path = Path(gguf_file)
-    quant_dir_re = (
-        r"(UD-)?("
-        r"MXFP[0-9]+(?:_[A-Z0-9]+)*"
-        r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?"
-        r"|TQ[0-9]+_[0-9]+"
-        r"|Q[0-9]+_K_[A-Z]+"
-        r"|Q[0-9]+_[0-9]+"
-        r"|Q[0-9]+_K"
-        r"|BF16|F16|F32"
-        r")"
-    )
+    # One quant vocabulary, shared: a local copy of it silently fell behind on
+    # the bpw modifier, which left IQ4_XS-3.53bpw unrecognised as a quant dir.
+    quant_dir_re = rf"{_GGUF_KNOWN_QUANT_RE.pattern}(-[0-9]+(?:\.[0-9]+)?bpw)?"
     search_dir = gguf_path.parent if selected.suffix.lower() == ".gguf" else selected
     if not search_dir.name:
         return str(search_dir)