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)